Repository: quora/pyanalyze
Pull request #645 (2023-06-12 18:04:29)

Add disallowed_imports


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 5ec06f9..663cbde 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Add `disallowed_imports` configuration option to disallow
+  imports of specific modules (#645)
 - Consider an annotated assignment without a value to be
   an exported name (#644)
 - Improve the location where `missing_parameter_annotation`
diff --git a/pyanalyze/error_code.py b/pyanalyze/error_code.py
index 5351b9f..9a0def3 100644
--- a/pyanalyze/error_code.py
+++ b/pyanalyze/error_code.py
@@ -106,6 +106,7 @@ class ErrorCode(enum.Enum):
     override_does_not_override = 86
     reveal_type = 87
     missing_generic_parameters = 88
+    disallowed_import = 89
 
 
 # Allow testing unannotated functions without too much fuss
@@ -237,6 +238,7 @@ ERROR_DESCRIPTION = {
     ErrorCode.invalid_override_decorator: "@override decorator in invalid location",
     ErrorCode.override_does_not_override: "Method does not override any base method",
     ErrorCode.missing_generic_parameters: "Missing type parameters for generic type",
+    ErrorCode.disallowed_import: "Disallowed import",
 }
 
 
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index 43d577b..7d48410 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -426,6 +426,16 @@ class DisallowCallsToDunders(StringSequenceOption):
     name = "disallow_calls_to_dunders"
 
 
+class DisallowedImports(StringSequenceOption):
+    """List of imports that will trigger an error.
+
+    Entries may be top-level modules (e.g., "os") or dotted submodule paths (e.g., "os.path").
+
+    """
+
+    name = "disallowed_imports"
+
+
 class ForLoopAlwaysEntered(BooleanOption):
     """If True, we assume that for loops are always entered at least once,
     which affects the potentially_undefined_name check. This will miss
@@ -2310,6 +2320,16 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
 
     # Imports
 
+    def check_for_disallowed_import(self, node: ast.AST, name: str) -> None:
+        print("CHECK", name)
+        disallowed = self.options.get_value_for(DisallowedImports)
+        if name in disallowed:
+            self._show_error_if_checking(
+                node,
+                f"Disallowed import of module {name!r}",
+                error_code=ErrorCode.disallowed_import,
+            )
+
     def visit_Import(self, node: ast.Import) -> None:
         self.generic_visit(node)
         if self.scopes.scope_type() == ScopeType.module_scope:
@@ -2317,6 +2337,7 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 self.import_name_to_node[name.name] = node
 
         for alias in node.names:
+            self.check_for_disallowed_import(node, alias.name)
             self._try_to_import(alias.name)
             # "import a.b" sets the name "a", but "import a.b as c" sets "c" to the value "a.b"
             varname = (
@@ -2404,6 +2425,18 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
             for name in node.names:
                 self.future_imports.add(name.name)
 
+        if node.module is not None and node.level == 0:
+            self.check_for_disallowed_import(node, node.module)
+            for alias in node.names:
+                # Before 3.10 the alias doesn't have a lineno
+                if sys.version_info >= (3, 10):
+                    error_node = alias
+                else:
+                    error_node = node
+                self.check_for_disallowed_import(
+                    error_node, f"{node.module}.{alias.name}"
+                )
+
         self._maybe_record_usages_from_import(node)
 
         # See if we can get the names from the stub instead


Test patch:
diff --git a/pyanalyze/test.toml b/pyanalyze/test.toml
index e551d77..5132a1b 100644
--- a/pyanalyze/test.toml
+++ b/pyanalyze/test.toml
@@ -12,4 +12,8 @@ functions_safe_to_call = [
 ]
 class_attribute_transformers = [
     "pyanalyze.test_config.transform_class_attribute"
+]
+disallowed_imports = [
+    "getopt",
+    "email.quoprimime",
 ]
\ No newline at end of file
diff --git a/pyanalyze/test_import.py b/pyanalyze/test_import.py
index 7b01179..16a10fe 100644
--- a/pyanalyze/test_import.py
+++ b/pyanalyze/test_import.py
@@ -41,3 +41,44 @@ class TestImport(TestNameCheckVisitorBase):
                 assert_is_value(extensions, KnownValue(P.extensions))
                 not_a_name  # E: undefined_name
             """)
+
+
+class TestDisallowedImport(TestNameCheckVisitorBase):
+    @assert_passes()
+    def test_top_level(self):
+        import getopt  # E: disallowed_import
+
+        from getopt import GetoptError  # E: disallowed_import
+
+        print(getopt, GetoptError)  # shut up flake8
+
+        def capybara():
+            import getopt  # E: disallowed_import
+            from getopt import GetoptError  # E: disallowed_import
+
+            print(getopt, GetoptError)
+
+    @assert_passes()
+    def test_nested(self):
+        import email.quoprimime  # E: disallowed_import
+        from email.quoprimime import unquote  # E: disallowed_import
+
+        print(email, unquote)
+
+        def capybara():
+            import email.quoprimime  # E: disallowed_import
+            from email.quoprimime import unquote  # E: disallowed_import
+            from email import quoprimime  # E: disallowed_import
+
+            print(email, unquote, quoprimime)
+
+    @assert_passes()
+    def test_import_from(self):
+        from email import quoprimime  # E: disallowed_import
+
+        print(quoprimime)
+
+        def capybara():
+            from email import quoprimime  # E: disallowed_import
+
+            print(quoprimime)