Modify visibility and shuffle around some modules (#1807)

src/source_code/locator.rs

@@ -0,0 +1,65 @@
+//! Struct used to efficiently slice source code at (row, column) Locations.
+
+use std::borrow::Cow;
+
+use once_cell::unsync::OnceCell;
+use ropey::Rope;
+use rustpython_ast::Location;
+
+use crate::ast::types::Range;
+
+pub struct Locator<'a> {
+    contents: &'a str,
+    rope: OnceCell<Rope>,
+}
+
+impl<'a> Locator<'a> {
+    pub fn new(contents: &'a str) -> Self {
+        Locator {
+            contents,
+            rope: OnceCell::default(),
+        }
+    }
+
+    fn get_or_init_rope(&self) -> &Rope {
+        self.rope.get_or_init(|| Rope::from_str(self.contents))
+    }
+
+    pub fn slice_source_code_at(&self, location: &Location) -> Cow<'_, str> {
+        let rope = self.get_or_init_rope();
+        let offset = rope.line_to_char(location.row() - 1) + location.column();
+        Cow::from(rope.slice(offset..))
+    }
+
+    pub fn slice_source_code_until(&self, location: &Location) -> Cow<'_, str> {
+        let rope = self.get_or_init_rope();
+        let offset = rope.line_to_char(location.row() - 1) + location.column();
+        Cow::from(rope.slice(..offset))
+    }
+
+    pub fn slice_source_code_range(&self, range: &Range) -> Cow<'_, str> {
+        let rope = self.get_or_init_rope();
+        let start = rope.line_to_char(range.location.row() - 1) + range.location.column();
+        let end = rope.line_to_char(range.end_location.row() - 1) + range.end_location.column();
+        Cow::from(rope.slice(start..end))
+    }
+
+    pub fn partition_source_code_at(
+        &self,
+        outer: &Range,
+        inner: &Range,
+    ) -> (Cow<'_, str>, Cow<'_, str>, Cow<'_, str>) {
+        let rope = self.get_or_init_rope();
+        let outer_start = rope.line_to_char(outer.location.row() - 1) + outer.location.column();
+        let outer_end =
+            rope.line_to_char(outer.end_location.row() - 1) + outer.end_location.column();
+        let inner_start = rope.line_to_char(inner.location.row() - 1) + inner.location.column();
+        let inner_end =
+            rope.line_to_char(inner.end_location.row() - 1) + inner.end_location.column();
+        (
+            Cow::from(rope.slice(outer_start..inner_start)),
+            Cow::from(rope.slice(inner_start..inner_end)),
+            Cow::from(rope.slice(inner_end..outer_end)),
+        )
+    }
+}

src/source_code/mod.rs

@@ -0,0 +1,7 @@
+mod generator;
+mod locator;
+mod stylist;
+
+pub use generator::Generator;
+pub use locator::Locator;
+pub use stylist::{LineEnding, Stylist};

src/source_code/stylist.rs

@@ -0,0 +1,279 @@
+//! Detect code style from Python source code.
+
+use std::fmt;
+use std::ops::Deref;
+
+use once_cell::unsync::OnceCell;
+use rustpython_ast::Location;
+use rustpython_parser::lexer;
+use rustpython_parser::lexer::Tok;
+
+use crate::ast::types::Range;
+use crate::pydocstyle::helpers::leading_quote;
+use crate::source_code::Locator;
+use crate::vendor;
+
+pub struct Stylist<'a> {
+    contents: &'a str,
+    locator: &'a Locator<'a>,
+    indentation: OnceCell<Indentation>,
+    quote: OnceCell<Quote>,
+    line_ending: OnceCell<LineEnding>,
+}
+
+impl<'a> Stylist<'a> {
+    pub fn indentation(&'a self) -> &'a Indentation {
+        self.indentation
+            .get_or_init(|| detect_indentation(self.contents, self.locator).unwrap_or_default())
+    }
+
+    pub fn quote(&'a self) -> &'a Quote {
+        self.quote
+            .get_or_init(|| detect_quote(self.contents, self.locator).unwrap_or_default())
+    }
+
+    pub fn line_ending(&'a self) -> &'a LineEnding {
+        self.line_ending
+            .get_or_init(|| detect_line_ending(self.contents).unwrap_or_default())
+    }
+
+    pub fn from_contents(contents: &'a str, locator: &'a Locator<'a>) -> Self {
+        Self {
+            contents,
+            locator,
+            indentation: OnceCell::default(),
+            quote: OnceCell::default(),
+            line_ending: OnceCell::default(),
+        }
+    }
+}
+
+/// The quotation style used in Python source code.
+#[derive(Debug, PartialEq, Eq)]
+pub enum Quote {
+    Single,
+    Double,
+}
+
+impl Default for Quote {
+    fn default() -> Self {
+        Quote::Double
+    }
+}
+
+impl From<&Quote> for vendor::str::Quote {
+    fn from(val: &Quote) -> Self {
+        match val {
+            Quote::Single => vendor::str::Quote::Single,
+            Quote::Double => vendor::str::Quote::Double,
+        }
+    }
+}
+
+impl fmt::Display for Quote {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Quote::Single => write!(f, "\'"),
+            Quote::Double => write!(f, "\""),
+        }
+    }
+}
+
+impl From<&Quote> for char {
+    fn from(val: &Quote) -> Self {
+        match val {
+            Quote::Single => '\'',
+            Quote::Double => '"',
+        }
+    }
+}
+
+/// The indentation style used in Python source code.
+#[derive(Debug, PartialEq, Eq)]
+pub struct Indentation(String);
+
+impl Indentation {
+    pub fn new(indentation: String) -> Self {
+        Self(indentation)
+    }
+}
+
+impl Default for Indentation {
+    fn default() -> Self {
+        Indentation("    ".to_string())
+    }
+}
+
+impl Deref for Indentation {
+    type Target = str;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+/// The line ending style used in Python source code.
+/// See <https://docs.python.org/3/reference/lexical_analysis.html#physical-lines>
+#[derive(Debug, PartialEq, Eq)]
+pub enum LineEnding {
+    Lf,
+    Cr,
+    CrLf,
+}
+
+impl Default for LineEnding {
+    fn default() -> Self {
+        LineEnding::Lf
+    }
+}
+
+impl Deref for LineEnding {
+    type Target = str;
+
+    fn deref(&self) -> &Self::Target {
+        match &self {
+            LineEnding::CrLf => "\r\n",
+            LineEnding::Lf => "\n",
+            LineEnding::Cr => "\r",
+        }
+    }
+}
+
+/// Detect the indentation style of the given tokens.
+fn detect_indentation(contents: &str, locator: &Locator) -> Option<Indentation> {
+    for (_start, tok, end) in lexer::make_tokenizer(contents).flatten() {
+        if let Tok::Indent { .. } = tok {
+            let start = Location::new(end.row(), 0);
+            let whitespace = locator.slice_source_code_range(&Range::new(start, end));
+            return Some(Indentation(whitespace.to_string()));
+        }
+    }
+    None
+}
+
+/// Detect the quotation style of the given tokens.
+fn detect_quote(contents: &str, locator: &Locator) -> Option<Quote> {
+    for (start, tok, end) in lexer::make_tokenizer(contents).flatten() {
+        if let Tok::String { .. } = tok {
+            let content = locator.slice_source_code_range(&Range::new(start, end));
+            if let Some(pattern) = leading_quote(&content) {
+                if pattern.contains('\'') {
+                    return Some(Quote::Single);
+                } else if pattern.contains('"') {
+                    return Some(Quote::Double);
+                }
+                unreachable!("Expected string to start with a valid quote prefix")
+            }
+        }
+    }
+    None
+}
+
+/// Detect the line ending style of the given contents.
+fn detect_line_ending(contents: &str) -> Option<LineEnding> {
+    if let Some(position) = contents.find('\n') {
+        let position = position.saturating_sub(1);
+        return if let Some('\r') = contents.chars().nth(position) {
+            Some(LineEnding::CrLf)
+        } else {
+            Some(LineEnding::Lf)
+        };
+    } else if contents.find('\r').is_some() {
+        return Some(LineEnding::Cr);
+    }
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::source_code::stylist::{
+        detect_indentation, detect_line_ending, detect_quote, Indentation, LineEnding, Quote,
+    };
+    use crate::source_code::Locator;
+
+    #[test]
+    fn indentation() {
+        let contents = r#"x = 1"#;
+        let locator = Locator::new(contents);
+        assert_eq!(detect_indentation(contents, &locator), None);
+
+        let contents = r#"
+if True:
+  pass
+"#;
+        let locator = Locator::new(contents);
+        assert_eq!(
+            detect_indentation(contents, &locator),
+            Some(Indentation("  ".to_string()))
+        );
+
+        let contents = r#"
+if True:
+    pass
+"#;
+        let locator = Locator::new(contents);
+        assert_eq!(
+            detect_indentation(contents, &locator),
+            Some(Indentation("    ".to_string()))
+        );
+
+        let contents = r#"
+if True:
+	pass
+"#;
+        let locator = Locator::new(contents);
+        assert_eq!(
+            detect_indentation(contents, &locator),
+            Some(Indentation("\t".to_string()))
+        );
+
+        // TODO(charlie): Should non-significant whitespace be detected?
+        let contents = r#"
+x = (
+  1,
+  2,
+  3,
+)
+"#;
+        let locator = Locator::new(contents);
+        assert_eq!(detect_indentation(contents, &locator), None);
+    }
+
+    #[test]
+    fn quote() {
+        let contents = r#"x = 1"#;
+        let locator = Locator::new(contents);
+        assert_eq!(detect_quote(contents, &locator), None);
+
+        let contents = r#"x = '1'"#;
+        let locator = Locator::new(contents);
+        assert_eq!(detect_quote(contents, &locator), Some(Quote::Single));
+
+        let contents = r#"x = "1""#;
+        let locator = Locator::new(contents);
+        assert_eq!(detect_quote(contents, &locator), Some(Quote::Double));
+
+        let contents = r#"
+def f():
+    """Docstring."""
+    pass
+"#;
+        let locator = Locator::new(contents);
+        assert_eq!(detect_quote(contents, &locator), Some(Quote::Double));
+    }
+
+    #[test]
+    fn line_ending() {
+        let contents = "x = 1";
+        assert_eq!(detect_line_ending(contents), None);
+
+        let contents = "x = 1\n";
+        assert_eq!(detect_line_ending(contents), Some(LineEnding::Lf));
+
+        let contents = "x = 1\r";
+        assert_eq!(detect_line_ending(contents), Some(LineEnding::Cr));
+
+        let contents = "x = 1\r\n";
+        assert_eq!(detect_line_ending(contents), Some(LineEnding::CrLf));
+    }
+}