Do not remove imports via airflint (#10)

Instead we advise users to use autoflake to remove imports.
feluelle · Jun 20, 2022 · cff2063 · cff2063
1 parent bec024f
commit cff2063
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -13,11 +13,13 @@
 
 ## 🧑‍🏫 Rules
 
-- [x] Use function-level imports instead of top-level imports[^1] (see [Top level Python Code](https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#top-level-python-code))
+- [x] Use function-level imports instead of top-level imports[^1][^2] (see [Top level Python Code](https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#top-level-python-code))
 - [x] Use jinja macro instead of `Variable.get` (see [Airflow Variables](https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#airflow-variables))
 
 [^1]: There is a PEP for [Lazy Imports](https://peps.python.org/pep-0690/) targeted to arrive in Python 3.12 which would supersede this rule.
 
+[^2]: To remove top-level imports after running `UseFunctionLevelImports` rule, use a tool such as [autoflake](https://github.com/PyCQA/autoflake).
+
 _based on official [Best Practices](https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html)_
 
 ## 🚀 Get started
@@ -43,3 +45,15 @@ Alternatively you can add the following repo to your `pre-commit-config.yaml`:
       - id: airflint
         args: ["-a"]  # Use -a for replacing inplace
 ```
+
+To complete the `UseFunctionlevelImports` rule, please add the `autoflake` hook after the `airflint` hook, as below:
+
+```yaml
+  - repo: https://github.com/pycqa/autoflake
+    rev: v1.4
+    hooks:
+      - id: autoflake
+        args: ["--remove-all-unused-imports", "--in-place"]
+```
+
+This will remove unused imports.
diff --git a/airflint/actions/move_statements.py → airflint/actions/new_statements.py b/airflint/actions/move_statements.py → airflint/actions/new_statements.py
@@ -7,8 +7,8 @@
 
 
 @dataclass
-class MoveStatementsAction(Action):
-    """Move statements from their old location to the node's location."""
+class NewStatementsAction(Action):
+    """Add new statements after the node's line."""
 
     statements: list[ast.stmt]
 
@@ -46,9 +46,5 @@ def apply(self, context: Context, source: str) -> str:
         for line in reversed(replacement):
             lines.insert(anchor, line)
 
-        # We remove the old statements.
-        for line in split_lines(new_source):
-            lines.remove(line)
-
         # Finally we'll merge everything together!
         return lines.join()
diff --git a/airflint/rules/use_function_level_imports.py b/airflint/rules/use_function_level_imports.py
@@ -4,7 +4,7 @@
 from refactor import Action, Rule, context
 from refactor.context import ScopeType
 
-from airflint.actions.move_statements import MoveStatementsAction
+from airflint.actions.new_statements import NewStatementsAction
 
 
 class UseFunctionLevelImports(Rule):
@@ -56,4 +56,4 @@ def match(self, node: ast.AST) -> Action:
         # We'll select the first statement, which will act like an anchor to us
         # when we are inserting.
         first_stmt = node.body[0]
-        return MoveStatementsAction(first_stmt, inlined_imports)
+        return NewStatementsAction(first_stmt, inlined_imports)
diff --git a/tests/test_rules.py b/tests/test_rules.py
@@ -26,6 +26,9 @@ def other_thing():
                 return dataclass(something(1, 2))
             """,
             """
+            import functools
+            import operator
+            import dataclass
 
             def something():
                 import functools