Fix splice

pyproject.toml

@@ -42,9 +42,3 @@ license-files = [
   "LICENSE",
   "licenses/*",
 ]
-
-[tool.isort]
-add_imports = "from __future__ import annotations"
-
-[tool.ruff.isort]
-required-imports = ["from __future__ import annotations"]

resources/test/fixtures/isort/required_imports/comment.py

@@ -0,0 +1,3 @@
+#!/usr/bin/env python3
+
+x = 1

resources/test/fixtures/isort/required_imports/docstring.py

@@ -0,0 +1,3 @@
+"""Hello, world!"""
+
+x = 1

resources/test/fixtures/isort/required_imports/docstring_only.py

@@ -0,0 +1 @@
+"""Hello, world!"""

resources/test/fixtures/isort/required_imports/docstring_with_continuation.py

@@ -0,0 +1,2 @@
+"""Hello, world!"""; x = \
+    1; y = 2

resources/test/fixtures/isort/required_imports/docstring_with_semicolon.py

@@ -0,0 +1 @@
+"""Hello, world!"""; x = 1

resources/test/fixtures/isort/required_imports/existing_import.py

@@ -0,0 +1,2 @@
+from __future__ import generator_stop
+import os

src/ast/helpers.rs

@@ -714,6 +714,21 @@ pub fn followed_by_multi_statement_line(stmt: &Stmt, locator: &SourceCodeLocator
     match_trailing_content(stmt, locator)
 }
 
+/// Return `true` if a `Stmt` is a docstring.
+pub fn is_docstring_stmt(stmt: &Stmt) -> bool {
+    if let StmtKind::Expr { value } = &stmt.node {
+        matches!(
+            value.node,
+            ExprKind::Constant {
+                value: Constant::Str { .. },
+                ..
+            }
+        )
+    } else {
+        false
+    }
+}
+
 #[derive(Default)]
 /// A simple representation of a call's positional and keyword arguments.
 pub struct SimpleCallArgs<'a> {

src/checkers/imports.rs

@@ -15,7 +15,6 @@ use crate::source_code_style::SourceCodeStyleDetector;
 
 #[allow(clippy::too_many_arguments)]
 pub fn check_imports(
-    contents: &str,
     python_ast: &Suite,
     locator: &SourceCodeLocator,
     directives: &IsortDirectives,
@@ -25,11 +24,6 @@ pub fn check_imports(
     path: &Path,
     package: Option<&Path>,
 ) -> Vec<Diagnostic> {
-    // Don't enforce import rules on empty files (like `__init__.py`).
-    if contents.is_empty() {
-        return vec![];
-    }
-
     // Extract all imports from the AST.
     let tracker = {
         let mut tracker = ImportTracker::new(locator, directives, path);
@@ -55,7 +49,7 @@ pub fn check_imports(
     }
     if settings.enabled.contains(&RuleCode::I002) {
         diagnostics.extend(isort::rules::add_required_imports(
-            contents, &blocks, settings, autofix,
+            &blocks, python_ast, locator, settings, autofix,
         ));
     }
     diagnostics

src/flake8_unused_arguments/helpers.rs

@@ -1,19 +1,6 @@
 use rustpython_ast::{Constant, ExprKind, Stmt, StmtKind};
 
-/// Return `true` if a `Stmt` is a docstring.
-fn is_docstring_stmt(stmt: &Stmt) -> bool {
-    if let StmtKind::Expr { value } = &stmt.node {
-        matches!(
-            value.node,
-            ExprKind::Constant {
-                value: Constant::Str { .. },
-                ..
-            }
-        )
-    } else {
-        false
-    }
-}
+use crate::ast::helpers::is_docstring_stmt;
 
 /// Return `true` if a `Stmt` is a "empty": a `pass`, `...`, `raise
 /// NotImplementedError`, or `raise NotImplemented` (with or without arguments).

src/isort/helpers.rs

@@ -1,7 +1,8 @@
-use rustpython_ast::Stmt;
+use rustpython_ast::{Location, Stmt};
 use rustpython_parser::lexer;
 use rustpython_parser::lexer::Tok;
 
+use crate::ast::helpers::is_docstring_stmt;
 use crate::ast::types::Range;
 use crate::isort::types::TrailingComma;
 use crate::source_code_locator::SourceCodeLocator;
@@ -86,3 +87,122 @@ pub fn has_comment_break(stmt: &Stmt, locator: &SourceCodeLocator) -> bool {
     }
     false
 }
+
+/// Find the end of the last docstring.
+fn match_docstring_end(body: &[Stmt]) -> Option<Location> {
+    let mut iter = body.iter();
+    let Some(mut stmt) = iter.next() else {
+        return None;
+    };
+    if !is_docstring_stmt(stmt) {
+        return None;
+    }
+    for next in iter {
+        if !is_docstring_stmt(next) {
+            break;
+        }
+        stmt = next;
+    }
+    Some(stmt.end_location.unwrap())
+}
+
+/// Find the end of the first token that isn't a docstring, comment, or
+/// whitespace.
+pub fn find_splice_location(body: &[Stmt], locator: &SourceCodeLocator) -> Location {
+    // Find the first AST node that isn't a docstring.
+    let mut splice = match_docstring_end(body).unwrap_or_default();
+
+    // Find the first token that isn't a comment or whitespace.
+    let contents = locator.slice_source_code_at(&splice);
+    for (.., tok, end) in lexer::make_tokenizer(&contents).flatten() {
+        if matches!(tok, Tok::Comment(..) | Tok::Newline) {
+            splice = end;
+        } else {
+            break;
+        }
+    }
+
+    splice
+}
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use rustpython_ast::Location;
+    use rustpython_parser::parser;
+
+    use crate::isort::helpers::find_splice_location;
+    use crate::source_code_locator::SourceCodeLocator;
+
+    fn splice_contents(contents: &str) -> Result<Location> {
+        let program = parser::parse_program(contents, "<filename>")?;
+        let locator = SourceCodeLocator::new(contents);
+        Ok(find_splice_location(&program, &locator))
+    }
+
+    #[test]
+    fn splice() -> Result<()> {
+        let contents = "";
+        assert_eq!(splice_contents(contents)?, Location::new(1, 0));
+
+        let contents = r#"
+"""Hello, world!"""
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(1, 19));
+
+        let contents = r#"
+"""Hello, world!"""
+"""Hello, world!"""
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(2, 19));
+
+        let contents = r#"
+x = 1
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(1, 0));
+
+        let contents = r#"
+#!/usr/bin/env python3
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(1, 22));
+
+        let contents = r#"
+#!/usr/bin/env python3
+"""Hello, world!"""
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(2, 19));
+
+        let contents = r#"
+"""Hello, world!"""
+#!/usr/bin/env python3
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(2, 22));
+
+        let contents = r#"
+"""%s""" % "Hello, world!"
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(1, 0));
+
+        let contents = r#"
+"""Hello, world!"""; x = 1
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(1, 19));
+
+        let contents = r#"
+"""Hello, world!"""; x = 1; y = \
+    2
+"#
+        .trim();
+        assert_eq!(splice_contents(contents)?, Location::new(1, 19));
+
+        Ok(())
+    }
+}

src/isort/mod.rs

@@ -828,4 +828,99 @@ mod tests {
         insta::assert_yaml_snapshot!(snapshot, diagnostics);
         Ok(())
     }
+
+    #[test_case(Path::new("docstring.py"))]
+    #[test_case(Path::new("docstring_only.py"))]
+    #[test_case(Path::new("empty.py"))]
+    fn required_import(path: &Path) -> Result<()> {
+        let snapshot = format!("required_import_{}", path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("./resources/test/fixtures/isort/required_imports")
+                .join(path)
+                .as_path(),
+            &Settings {
+                src: vec![Path::new("resources/test/fixtures/isort").to_path_buf()],
+                isort: isort::settings::Settings {
+                    required_imports: BTreeSet::from([
+                        "from __future__ import annotations".to_string()
+                    ]),
+                    ..isort::settings::Settings::default()
+                },
+                ..Settings::for_rule(RuleCode::I002)
+            },
+        )?;
+        insta::assert_yaml_snapshot!(snapshot, diagnostics);
+        Ok(())
+    }
+
+    #[test_case(Path::new("docstring.py"))]
+    #[test_case(Path::new("docstring_only.py"))]
+    #[test_case(Path::new("empty.py"))]
+    fn required_imports(path: &Path) -> Result<()> {
+        let snapshot = format!("required_imports_{}", path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("./resources/test/fixtures/isort/required_imports")
+                .join(path)
+                .as_path(),
+            &Settings {
+                src: vec![Path::new("resources/test/fixtures/isort").to_path_buf()],
+                isort: isort::settings::Settings {
+                    required_imports: BTreeSet::from([
+                        "from __future__ import annotations".to_string(),
+                        "from __future__ import generator_stop".to_string(),
+                    ]),
+                    ..isort::settings::Settings::default()
+                },
+                ..Settings::for_rule(RuleCode::I002)
+            },
+        )?;
+        insta::assert_yaml_snapshot!(snapshot, diagnostics);
+        Ok(())
+    }
+
+    #[test_case(Path::new("docstring.py"))]
+    #[test_case(Path::new("docstring_only.py"))]
+    #[test_case(Path::new("empty.py"))]
+    fn combined_required_imports(path: &Path) -> Result<()> {
+        let snapshot = format!("combined_required_imports_{}", path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("./resources/test/fixtures/isort/required_imports")
+                .join(path)
+                .as_path(),
+            &Settings {
+                src: vec![Path::new("resources/test/fixtures/isort").to_path_buf()],
+                isort: isort::settings::Settings {
+                    required_imports: BTreeSet::from(["from __future__ import annotations, \
+                                                       generator_stop"
+                        .to_string()]),
+                    ..isort::settings::Settings::default()
+                },
+                ..Settings::for_rule(RuleCode::I002)
+            },
+        )?;
+        insta::assert_yaml_snapshot!(snapshot, diagnostics);
+        Ok(())
+    }
+
+    #[test_case(Path::new("docstring.py"))]
+    #[test_case(Path::new("docstring_only.py"))]
+    #[test_case(Path::new("empty.py"))]
+    fn straight_required_import(path: &Path) -> Result<()> {
+        let snapshot = format!("straight_required_import_{}", path.to_string_lossy());
+        let diagnostics = test_path(
+            Path::new("./resources/test/fixtures/isort/required_imports")
+                .join(path)
+                .as_path(),
+            &Settings {
+                src: vec![Path::new("resources/test/fixtures/isort").to_path_buf()],
+                isort: isort::settings::Settings {
+                    required_imports: BTreeSet::from(["import os".to_string()]),
+                    ..isort::settings::Settings::default()
+                },
+                ..Settings::for_rule(RuleCode::I002)
+            },
+        )?;
+        insta::assert_yaml_snapshot!(snapshot, diagnostics);
+        Ok(())
+    }
 }