Enum pass through (#186)

* pass through enums as-is so the C++ layer can deal with them.

* Detect a special namespace addition when working with enums.

* Fix up pylance error

* Ignore the injected namespace.

* Also mark other known enums.

* Fix 3.9 support issue

Author: gordonwatts

func_adl/type_based_replacement.py

@@ -906,7 +906,8 @@ def visit_Name(self, node: ast.Name) -> ast.Name:
             elif node.id in _global_functions:
                 self._found_types[node] = Callable
             else:
-                logging.getLogger(__name__).warning(f"Unknown type for name {node.id}")
+                if not getattr(node, "_ignore", False):
+                    logging.getLogger(__name__).warning(f"Unknown type for name {node.id}")
                 self._found_types[node] = Any
             return node
 

func_adl/util_ast.py

@@ -1,13 +1,13 @@
-from __future__ import annotations
-
 import ast
+import copy
+import importlib
 import inspect
 import tokenize
 from collections import defaultdict
 from dataclasses import is_dataclass
 from enum import Enum
 from types import ModuleType
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union, cast
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 
 
 def as_literal(p: Union[str, int, float, bool, None]) -> ast.Constant:
@@ -252,6 +252,13 @@ def rewrite_func_as_lambda(f: ast.FunctionDef) -> ast.Lambda:
     return ast.Lambda(args, ret.value)  # type: ignore
 
 
+class _mark_ignore_name(ast.NodeTransformer):
+    def visit_Name(self, node: ast.Name) -> Any:
+        new_node = cast(ast.Expr, ast.parse(node.id).body[0]).value
+        new_node._ignore = True  # type: ignore
+        return new_node
+
+
 class _rewrite_captured_vars(ast.NodeTransformer):
     def __init__(self, cv: inspect.ClosureVars):
         self._lookup_dict: Dict[str, Any] = dict(cv.nonlocals)
@@ -293,7 +300,23 @@ def visit_Attribute(self, node: ast.Attribute) -> Any:
             new_value = getattr(value.value, node.attr)
             # When 3.10 is not supported, replace with EnumType
             if isinstance(value.value, Enum.__class__):
-                new_value = new_value.value
+                # Sometimes we need to prepend a namespace. We look
+                # for secret info here, and then prepend if necessary.
+                # But no one else knows about this, so  we need to mark this
+                # as "ignore".
+                enum_mod = importlib.import_module(new_value.__module__)
+                additional_ns = getattr(enum_mod, "_object_cpp_as_py_namespace", None)
+                if additional_ns is None:
+                    return node
+
+                if len(additional_ns) > 0:
+                    ns_node = cast(
+                        ast.Expr, ast.parse(f"{additional_ns}.{ast.unparse(node)}").body[0]
+                    ).value
+                else:
+                    ns_node = copy.copy(node)
+
+                return _mark_ignore_name().visit(ns_node)
             return ast.Constant(value=new_value)
 
         # If we fail, then just move on.
@@ -448,9 +471,7 @@ def find_identifier(
                 break
         return None, None
 
-    def tokens_till(
-        self, stop_condition: Dict[int, List[str]]
-    ) -> Generator[tokenize.Token, None, None]:  # type: ignore
+    def tokens_till(self, stop_condition: Dict[int, List[str]]):
         """Yield tokens until we find a stop condition.
 
         * Properly tracks parentheses, etc.

tests/test_type_based_replacement.py

@@ -3,7 +3,7 @@
 import inspect
 import logging
 from inspect import isclass
-from typing import Any, Callable, Iterable, Optional, Tuple, Type, TypeVar, cast
+from typing import Any, Callable, Iterable, Optional, Tuple, Type, TypeVar, Union, cast
 
 import pytest
 
@@ -128,8 +128,8 @@ def EventNumber(self) -> int: ...  # noqa
     def MyLambdaCallback(self, cb: Callable) -> int: ...  # noqa
 
 
-def return_type_test(expr: str, arg_type: type, expected_type: type):
-    s = ast_lambda(expr)
+def return_type_test(expr: Union[str, ast.expr], arg_type: type, expected_type: type):
+    s = expr if isinstance(expr, ast.expr) else ast_lambda(expr)
     objs = ObjectStream(ast.Name(id="e", ctx=ast.Load()), arg_type)
 
     _, _, expr_type = remap_by_types(objs, {"e": arg_type}, s)
@@ -251,6 +251,17 @@ def test_subscript_any():
     return_type_test("e[0]", Any, Any)  # type: ignore
 
 
+def test_ast_marked_ignore(caplog):
+    "Make sure that a ast marked ignore does not generate a warning"
+    a = ast.parse("donttalk").body[0].value  # type: ignore
+    a._ignore = True
+
+    with caplog.at_level(logging.WARNING):
+        return_type_test(a, Any, Any)  # type: ignore
+
+    assert len(caplog.messages) == 0
+
+
 def test_collection():
     "A simple collection"
     s = ast_lambda("e.Jets('default')")

tests/test_util_ast.py

@@ -292,9 +292,76 @@ class MyEnum(Enum):
             VALUE = 20
 
     r = parse_as_ast(lambda x: x > forkit.MyEnum.VALUE)
-    r_true = parse_as_ast(lambda x: x > 20)
 
-    assert ast.unparse(r) == ast.unparse(r_true)
+    assert "VALUE" in ast.unparse(r)
+
+
+def test_parse_lambda_with_implied_ns():
+    "Test adding the special attribute to the module to prefix a namespace"
+    # Add the attribute to the module
+    global _object_cpp_as_py_namespace
+    _object_cpp_as_py_namespace = "aweful"
+
+    try:
+
+        class forkit:
+            class MyEnum(Enum):
+                VALUE = 20
+
+        r = parse_as_ast(lambda x: x > forkit.MyEnum.VALUE)
+
+        assert "aweful.forkit.MyEnum.VALUE" in ast.unparse(r)
+
+        found_it = False
+
+        class check_it(ast.NodeVisitor):
+            def visit_Name(self, node: ast.Name):
+                nonlocal found_it
+                if node.id == "aweful":
+                    found_it = True
+                    assert hasattr(node, "_ignore")
+                    assert node._ignore  # type: ignore
+
+        check_it().visit(r)
+        assert found_it
+
+    finally:
+        # Remove the attribute from the module
+        del _object_cpp_as_py_namespace
+
+
+def test_parse_lambda_with_implied_ns_empty():
+    "Test adding the special attribute to the module to prefix a namespace"
+    # Add the attribute to the module
+    global _object_cpp_as_py_namespace
+    _object_cpp_as_py_namespace = ""
+
+    try:
+
+        class forkit:
+            class MyEnum(Enum):
+                VALUE = 20
+
+        r = parse_as_ast(lambda x: x > forkit.MyEnum.VALUE)
+
+        assert "forkit.MyEnum.VALUE" in ast.unparse(r)
+
+        found_it = False
+
+        class check_it(ast.NodeVisitor):
+            def visit_Name(self, node: ast.Name):
+                nonlocal found_it
+                if node.id == "forkit":
+                    found_it = True
+                    assert hasattr(node, "_ignore")
+                    assert node._ignore  # type: ignore
+
+        check_it().visit(r)
+        assert found_it
+
+    finally:
+        # Remove the attribute from the module
+        del _object_cpp_as_py_namespace
 
 
 def test_parse_lambda_class_constant_in_module():