1  from pathlib import Path
 2  
 3  from clint.config import Config
 4  from clint.index import SymbolIndex
 5  from clint.linter import Position, Range, lint_file
 6  from clint.rules.forbidden_top_level_import import ForbiddenTopLevelImport
 7  
 8  
 9  def test_forbidden_top_level_import(index: SymbolIndex) -> None:
10      code = """
11  # Bad
12  import foo
13  from foo import bar
14  
15  # Good
16  import baz
17  """
18      config = Config(
19          select={ForbiddenTopLevelImport.name},
20          forbidden_top_level_imports={"*": ["foo"]},
21      )
22      violations = lint_file(Path("test.py"), code, config, index)
23      assert len(violations) == 2
24      assert all(isinstance(v.rule, ForbiddenTopLevelImport) for v in violations)
25      assert violations[0].range == Range(Position(2, 0))
26      assert violations[1].range == Range(Position(3, 0))
27  
28  
29  def test_nested_if_in_type_checking_block(index: SymbolIndex) -> None:
30      code = """
31  from typing import TYPE_CHECKING
32  
33  if TYPE_CHECKING:
34      if True:
35          pass
36      import databricks  # Should NOT be flagged
37      from databricks import foo  # Should NOT be flagged
38  """
39      config = Config(
40          select={ForbiddenTopLevelImport.name},
41          forbidden_top_level_imports={"*": ["databricks"]},
42      )
43      violations = lint_file(Path("test.py"), code, config, index)
44      # Should have no violations since imports are inside TYPE_CHECKING
45      assert len(violations) == 0