cleanup

src/RestrictedPython/transformer.py

@@ -452,14 +452,7 @@ def check_import_names(self, node):
 
         This is a protection against rebinding dunder names like
         _getitem_, _write_ via imports.
-
-        => 'from _a import x' is ok, because '_a' is not added to the scope.
         """
-        if (isinstance(node, ast.ImportFrom)
-                and not node.module == '__future__'
-                and any(name.startswith('_') for name in node.module.split('.'))):  # NOQA: E501
-            self.error(node, 'module name starts "_", which is forbidden.')
-
         for name in node.names:
             if '*' in name.name:
                 self.error(node, '"*" imports are not allowed.')
@@ -1161,7 +1154,13 @@ def visit_Import(self, node):
 
     def visit_ImportFrom(self, node):
         """Allow `import from` statements with restrictions.
-        See check_import_names."""
+        See check_import_names.
+
+        => 'from _a import x' is ok, because '_a' is not added to the scope.
+        """
+        if not node.module == '__future__' and any(name.startswith('_') for name in node.module.split('.')):  # NOQA: E501
+            self.error(node, 'module name starts "_", which is forbidden.')
+
         return self.check_import_names(node)
 
     def visit_alias(self, node):

tests/transformer/test_import.py

@@ -46,25 +46,16 @@ def test_RestrictingNodeTransformer__visit_Import__5(c_exec):
 
 @pytest.mark.parametrize(*c_exec)
 def test_RestrictingNodeTransformer__visit_Import__6(c_exec):
-    """Deny imports from modules staring with `_`."""
+    """Deny imports from modules starting with `_`."""
+    message = 'Line 1: module name starts "_", which is forbidden.'
     result = c_exec('from _a import m')
-    assert result.errors == (
-        'Line 1: module name starts "_", which is forbidden.',
-    )
+    assert result.errors == (message,)
     result = c_exec('from a._b import m')
-    assert result.errors == (
-        'Line 1: module name starts "_", which is forbidden.',
-    )
+    assert result.errors == (message,)
     result = c_exec('from _a.b import m')
-    assert result.errors == (
-        'Line 1: module name starts "_", which is forbidden.',
-    )
+    assert result.errors == (message,)
     result = c_exec('from _a._b import m as x')
-    assert result.errors == (
-        'Line 1: module name starts "_", which is forbidden.',
-    )
-    result = c_exec('from a import m as _n')
-    assert result.errors == (import_errmsg % '_n',)
+    assert result.errors == (message,)
 
 
 @pytest.mark.parametrize(*c_exec)