feat(visitor): skip type checking blocks

src/visitor.rs

@@ -1,17 +1,19 @@
 use ruff_python_ast::visitor::{walk_stmt, Visitor};
-use ruff_python_ast::{self, Stmt};
+use ruff_python_ast::{self, Expr, ExprAttribute, ExprName, Stmt, StmtIf, StmtImportFrom};
 use ruff_text_size::TextRange;
 use std::collections::HashMap;
 
 #[derive(Debug, Clone)]
 pub struct ImportVisitor {
     imports: HashMap<String, Vec<TextRange>>,
+    has_future_annotations: bool,
 }
 
 impl ImportVisitor {
     pub fn new() -> Self {
         Self {
             imports: HashMap::new(),
+            has_future_annotations: false,
         }
     }
 
@@ -35,14 +37,20 @@ impl<'a> Visitor<'a> for ImportVisitor {
             Stmt::ImportFrom(import_from_stmt) => {
                 if let Some(module) = &import_from_stmt.module {
                     if import_from_stmt.level == Some(0) {
-                        // Assuming Int::new(0) is comparable with 0
+                        if is_future_annotations_import(module.as_str(), import_from_stmt) {
+                            self.has_future_annotations = true;
+                        }
+
                         self.imports
                             .entry(get_top_level_module_name(module.as_str()))
                             .or_default()
                             .push(import_from_stmt.range);
                     }
                 }
             }
+            Stmt::If(if_stmt) if is_typing_only_block(self.has_future_annotations, if_stmt) => {
+                // Avoid parsing imports that are only evaluated by type checkers.
+            }
             _ => walk_stmt(self, stmt), // Delegate other statements to walk_stmt
         }
     }
@@ -57,3 +65,37 @@ fn get_top_level_module_name(module_name: &str) -> String {
         .unwrap_or(module_name)
         .to_owned()
 }
+
+/// Checks if the import is a `from __future__ import annotations` one.
+fn is_future_annotations_import(module: &str, import_from_stmt: &StmtImportFrom) -> bool {
+    return module == "__future__"
+        && import_from_stmt
+            .names
+            .iter()
+            .any(|alias| alias.name.as_str() == "annotations");
+}
+
+/// Checks if we are in a block that will only be evaluated by type checkers, in accordance with
+/// <https://peps.python.org/pep-0563/>. If no `__future__.annotations` import is made, a block using `TYPE_CHECKING`
+/// will be evaluated at runtime, so we should not consider that this is a typing only block in that case.
+fn is_typing_only_block(has_future_annotations: bool, if_stmt: &StmtIf) -> bool {
+    if has_future_annotations {
+        match &if_stmt.test.as_ref() {
+            Expr::Attribute(ExprAttribute { value, attr, .. }) => {
+                if let Expr::Name(ExprName { id, .. }) = value.as_ref() {
+                    if id.as_str() == "typing" && attr.as_str() == "TYPE_CHECKING" {
+                        return true;
+                    }
+                }
+            }
+            Expr::Name(ExprName { id, .. }) => {
+                if id == "TYPE_CHECKING" {
+                    return true;
+                }
+            }
+            _ => (),
+        }
+    }
+
+    false
+}

tests/data/some_imports.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
+import typing
 from os import chdir, walk
 from pathlib import Path
-from typing import List
+from typing import List, TYPE_CHECKING
 
 import numpy as np
 import pandas
@@ -20,6 +23,12 @@
 import barfoo as bf
 from randomizer import random
 
+if TYPE_CHECKING:
+    import mypy_boto3_s3
+
+if typing.TYPE_CHECKING:
+    import mypy_boto3_sagemaker
+
 try:
     import click
 except:

tests/unit/imports/test_extract.py

@@ -21,23 +21,27 @@ def test_import_parser_py() -> None:
     some_imports_path = Path("tests/data/some_imports.py")
 
     assert get_imported_modules_from_list_of_files([some_imports_path]) == {
-        "barfoo": [Location(some_imports_path, 20, 8)],
-        "baz": [Location(some_imports_path, 16, 5)],
-        "click": [Location(some_imports_path, 24, 12)],
-        "foobar": [Location(some_imports_path, 18, 12)],
-        "httpx": [Location(some_imports_path, 14, 12)],
-        "module_in_class": [Location(some_imports_path, 35, 16)],
-        "module_in_func": [Location(some_imports_path, 30, 12)],
-        "not_click": [Location(some_imports_path, 26, 12)],
+        "__future__": [Location(some_imports_path, 1, 1)],
+        "barfoo": [Location(some_imports_path, 23, 8)],
+        "baz": [Location(some_imports_path, 19, 5)],
+        "click": [Location(some_imports_path, 33, 12)],
+        "foobar": [Location(some_imports_path, 21, 12)],
+        "httpx": [Location(some_imports_path, 17, 12)],
+        "module_in_class": [Location(some_imports_path, 44, 16)],
+        "module_in_func": [Location(some_imports_path, 39, 12)],
+        "not_click": [Location(some_imports_path, 35, 12)],
         "numpy": [
-            Location(some_imports_path, 5, 8),
-            Location(some_imports_path, 7, 1),
+            Location(some_imports_path, 8, 8),
+            Location(some_imports_path, 10, 1),
+        ],
+        "os": [Location(some_imports_path, 4, 1)],
+        "pandas": [Location(some_imports_path, 9, 8)],
+        "pathlib": [Location(some_imports_path, 5, 1)],
+        "randomizer": [Location(some_imports_path, 24, 1)],
+        "typing": [
+            Location(some_imports_path, 3, 8),
+            Location(some_imports_path, 6, 1),
         ],
-        "os": [Location(some_imports_path, 1, 1)],
-        "pandas": [Location(some_imports_path, 6, 8)],
-        "pathlib": [Location(some_imports_path, 2, 1)],
-        "randomizer": [Location(some_imports_path, 21, 1)],
-        "typing": [Location(some_imports_path, 3, 1)],
     }