test enum class

Author: gdraheim

tests/exectests3.py

@@ -60,18 +60,18 @@ def q_str(part: Union[str, int, None]) -> str:
     if isinstance(part, int):
         return str(part)
     return "'%s'" % part  # pylint: disable=consider-using-f-string
-def sh____(cmd: Union[str, List[str]], shell: bool = True) -> int:
+def sh____(cmd: Union[str, List[str]], shell: bool = True, env=None) -> int:
     if isinstance(cmd, string_types):
         logg.info(": %s", cmd)
     else:
         logg.info(": %s", " ".join([q_str(item) for item in cmd]))
-    return subprocess.check_call(cmd, shell=shell)
-def sx____(cmd: Union[str, List[str]], shell: bool = True) -> int:
+    return subprocess.check_call(cmd, shell=shell, env=env)
+def sx____(cmd: Union[str, List[str]], shell: bool = True, env=None) -> int:
     if isinstance(cmd, string_types):
         logg.info(": %s", cmd)
     else:
         logg.info(": %s", " ".join([q_str(item) for item in cmd]))
-    return subprocess.call(cmd, shell=shell)
+    return subprocess.call(cmd, shell=shell, env=env)
 
 class CalledProcessError(subprocess.SubprocessError):
     def __init__(self, args: Union[str, List[str]], returncode: int = 0, stdout: Union[str,bytes] = NIX, stderr: Union[str,bytes] = NIX) -> None:
@@ -889,6 +889,36 @@ class X(TypedDict):
             self.assertEqual(x2.out, "Success: no issues found in 1 source file")
         self.rm_testdir()
         self.end()
+    def test_3541(self) -> None:
+        """ check Enum classes are replaced by local def"""
+        vv = self.begin()
+        python = PYTHON
+        tmp = self.testdir()
+        text_file(F"{tmp}/test3.py", """
+        from enum import Enum
+        class A(Enum):
+            B = 2
+            C = 3
+        def func1() -> A:
+            return A(3)
+        f = func1()
+        if f is A.C:
+            print("OK:", f)
+        else:
+            print("NO:", f)
+        """)
+        sh____(F"{PYTHON3} {STRIP} -3 {tmp}/test3.py {vv}", env={"PYTHON3_ENUM_CLASS_ATLEAST":"3.99"})
+        self.assertTrue(os.path.exists(F"{tmp}/test.py"))
+        self.assertTrue(os.path.exists(F"{tmp}/test.pyi"))
+        script = lines4(open(F"{tmp}/test.py").read())
+        logg.info("script = %s", script)
+        self.assertTrue(greps(script, "(3, 99)"))
+        self.assertTrue(greps(script, "class Enum"))
+        x1 = X(F"{python} {tmp}/test.py")
+        logg.info("%s -> %s\n%s", x1.args, x1.out, x1.err)
+        self.assertTrue(greps(x1.out, "OK: <C: 3>"))
+        self.rm_testdir()
+        self.end()
 
 
 if __name__ == "__main__":

tests/transformertests2.py

@@ -4522,6 +4522,7 @@ def func1() -> int:
         self.coverage()
         self.rm_testdir()
     def test_2541(self) -> None:
+        """ check Enum classes are replaced by local def"""
         vv = self.begin()
         strip = coverage(STRIP)
         tmp = self.testdir()
@@ -4541,6 +4542,7 @@ def func1() -> int:
         py, pyi = file_text4(F"{tmp}/test.py"), file_text4(F"{tmp}/test.pyi")
         logg.debug("py:\n%s", py)
         self.assertEqual(lines4(py), lines4(text4("""
+        import sys
         if sys.version_info >= (3, 3):
             from enum import Enum
         else:

tool/strip_python3.py

@@ -1191,8 +1191,7 @@ def visit(self, node: ast.AST) -> ast.AST:
 
 class EnumClassTransformer(DetectImportsTransformer):
     typedefs: List[ast.stmt]
-    requiresfrom: Set[str]
-    only: Set[str]
+    requires: List[str]
     _Enum = """class Enum:
 
         def __new__(cls, *values):
@@ -1215,21 +1214,15 @@ def __iter__(self):
                 elem = getattr(cls, name)
                 if isinstance(elem, Enum):
                     yield elem
+
         def __str__(self):
             return "<%s: %s>" % (self.name, self.value)
     """
     def __init__(self) -> None:
         DetectImportsTransformer.__init__(self)
-        self.only = set()
-    def visit(self, node: ast.AST) -> ast.AST:
-        if isinstance(node, ast.Module):
-            module = cast(ast.Module, node)  # type: ignore[redundant-cast]
-            for stmt in module.body:
-                if isinstance(stmt, ast.ClassDef):
-                    self.only.add(stmt.name) # only top-level class names
-        return cast(ast.AST, DetectImportsTransformer.visit(self, node))
+        self.requires = []
     def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST: # pylint: disable=invalid-name
-        atleast = (3, 3)
+        atleast = [int(val) for val in os.environ.get("PYTHON3_ENUM_CLASS_ATLEAST", "3.3").split(".")]
         imports: ast.ImportFrom = node
         if imports.module and imports.module == "enum":
             orelse: List[ast.stmt] = []
@@ -1245,10 +1238,12 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST: # pylint: disable=i
             testcompare = testbody.value
             python2 = ast.If(test=testcompare, body=[imports], orelse=orelse)
             python2 = copy_location(python2, imports)
+            if "sys" not in self.requires:
+                self.requires += ["sys"]
             return python2
         return node
     def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST: # pylint: disable=invalid-name
-        atleast = (3, 3)
+        atleast = [int(val) for val in os.environ.get("PYTHON3_ENUM_CLASS_ATLEAST", "3.3").split(".")]
         classname = node.name
         for base in node.bases:
             if isinstance(base, ast.Name):
@@ -2528,6 +2523,7 @@ def visit(self, tree: ast.AST) -> ast.AST:
         if want.replace_enum_class:
             typedenum = EnumClassTransformer()
             tree = typedenum.visit(tree)
+            importrequires.append(typedenum.requires)
         extracted = ExtractTypeHints()
         tree = extracted.visit(tree)
         self.typedefs.extend(extracted.typedefs)