Repository: pantsbuild/pants

Issue:
Rule visitor chokes on union types
              It doesn't work, you get:

```
pants.base.exceptions.RuleTypeError: /Users/benjyw/src/pants/pants-plugins/internal_plugins/test_lockfile_fixtures/rules.py:85: Could not resolve type for `pytest.to_pex_request` in module internal_plugins.test_lockfile_fixtures.rules.

Return type annotation required for `to_pex_request` in /Users/benjyw/src/pants/src/python/pants/backend/python/subsystems/python_tool_base.py:276
```

Even though `to_pex_request` does have a return type annotation.

_Originally posted by @benjyw in https://github.com/pantsbuild/pants/pull/18350#discussion_r1117229506_
            


Discussion:

Resolving patch:
diff --git a/src/python/pants/engine/internals/rule_visitor.py b/src/python/pants/engine/internals/rule_visitor.py
--- a/src/python/pants/engine/internals/rule_visitor.py
+++ b/src/python/pants/engine/internals/rule_visitor.py
@@ -9,7 +9,7 @@
 import sys
 from contextlib import contextmanager
 from functools import partial
-from typing import Any, Callable, Iterator, List
+from typing import Any, Callable, Iterator, List, get_type_hints
 
 import typing_extensions
 
@@ -21,11 +21,12 @@
     GetParseError,
     MultiGet,
 )
-from pants.util.backport import get_annotations
 from pants.util.memo import memoized
 from pants.util.strutil import softwrap
+from pants.util.typing import patch_forward_ref
 
 logger = logging.getLogger(__name__)
+patch_forward_ref()
 
 
 def _get_starting_indent(source: str) -> int:
@@ -87,17 +88,16 @@ def pop(self) -> None:
 
 
 def _lookup_annotation(obj: Any, attr: str) -> Any:
-    """Get type assocated with a particular attribute on object. This can get hairy, especially on
+    """Get type associated with a particular attribute on object. This can get hairy, especially on
     Python <3.10.
 
     https://docs.python.org/3/howto/annotations.html#accessing-the-annotations-dict-of-an-object-in-python-3-9-and-older
-    For this reason, we've copied the `inspect.get_annotations` method from CPython `main` branch.
     """
     if hasattr(obj, attr):
         return getattr(obj, attr)
     else:
         try:
-            return get_annotations(obj, eval_str=True).get(attr)
+            return get_type_hints(obj).get(attr)
         except (NameError, TypeError):
             return None
 
@@ -113,7 +113,7 @@ def _lookup_return_type(func: Callable, check: bool = False) -> Any:
         func_file = inspect.getsourcefile(func)
         func_line = func.__code__.co_firstlineno
         raise TypeError(
-            f"Return type annotation required for `{func.__name__}` in {func_file}:{func_line}"
+            f"Failed to look up return type hint for `{func.__name__}` in {func_file}:{func_line}"
         )
     return ret
 
diff --git a/src/python/pants/util/backport.py b/src/python/pants/util/backport.py
deleted file mode 100644
--- a/src/python/pants/util/backport.py
+++ /dev/null
@@ -1,115 +0,0 @@
-# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
-# Licensed under the Apache License, Version 2.0 (see LICENSE).
-#
-# Copyright exception for `get_annotations()` which is in the public domain.  No warranties.
-# Copied from https://github.com/python/cpython/blob/951303fd855838d47765dcd05471e14311dc9fdd/Lib/inspect.py
-
-import functools
-import sys
-import types
-
-
-def get_annotations(obj, *, globals=None, locals=None, eval_str=False):
-    """Compute the annotations dict for an object. obj may be a callable, class, or module. Passing
-    in an object of any other type raises TypeError. Returns a dict.  get_annotations() returns a
-    new dict every time it's called; calling it twice on the same object will return two different
-    but equivalent dicts. This function handles several details for you:
-
-      * If eval_str is true, values of type str will
-        be un-stringized using eval().  This is intended
-        for use with stringized annotations
-        ("from __future__ import annotations").
-      * If obj doesn't have an annotations dict, returns an
-        empty dict.  (Functions and methods always have an
-        annotations dict; classes, modules, and other types of
-        callables may not.)
-      * Ignores inherited annotations on classes.  If a class
-        doesn't have its own annotations dict, returns an empty dict.
-      * All accesses to object members and dict values are done
-        using getattr() and dict.get() for safety.
-      * Always, always, always returns a freshly-created dict.
-    eval_str controls whether or not values of type str are replaced
-    with the result of calling eval() on those values:
-      * If eval_str is true, eval() is called on values of type str.
-      * If eval_str is false (the default), values of type str are unchanged.
-    globals and locals are passed in to eval(); see the documentation
-    for eval() for more information.  If either globals or locals is
-    None, this function may replace that value with a context-specific
-    default, contingent on type(obj):
-      * If obj is a module, globals defaults to obj.__dict__.
-      * If obj is a class, globals defaults to
-        sys.modules[obj.__module__].__dict__ and locals
-        defaults to the obj class namespace.
-      * If obj is a callable, globals defaults to obj.__globals__,
-        although if obj is a wrapped function (using
-        functools.update_wrapper()) it is first unwrapped.
-    """
-    if isinstance(obj, type):
-        # class
-        obj_dict = getattr(obj, "__dict__", None)
-        if obj_dict and hasattr(obj_dict, "get"):
-            ann = obj_dict.get("__annotations__", None)
-            if isinstance(ann, types.GetSetDescriptorType):
-                ann = None
-        else:
-            ann = None
-
-        obj_globals = None
-        module_name = getattr(obj, "__module__", None)
-        if module_name:
-            module = sys.modules.get(module_name, None)
-            if module:
-                obj_globals = getattr(module, "__dict__", None)
-        obj_locals = dict(vars(obj))
-        unwrap = obj
-    elif isinstance(obj, types.ModuleType):
-        # module
-        ann = getattr(obj, "__annotations__", None)
-        obj_globals = getattr(obj, "__dict__")
-        obj_locals = None
-        unwrap = None
-    elif callable(obj):
-        # this includes types.Function, types.BuiltinFunctionType,
-        # types.BuiltinMethodType, functools.partial, functools.singledispatch,
-        # "class funclike" from Lib/test/test_inspect... on and on it goes.
-        ann = getattr(obj, "__annotations__", None)
-        obj_globals = getattr(obj, "__globals__", None)
-        obj_locals = None
-        unwrap = obj
-    else:
-        raise TypeError(f"{obj!r} is not a module, class, or callable.")
-
-    if ann is None:
-        return {}
-
-    if not isinstance(ann, dict):
-        raise ValueError(f"{obj!r}.__annotations__ is neither a dict nor None")
-
-    if not ann:
-        return {}
-
-    if not eval_str:
-        return dict(ann)
-
-    if unwrap is not None:
-        while True:
-            if hasattr(unwrap, "__wrapped__"):
-                unwrap = unwrap.__wrapped__
-                continue
-            if isinstance(unwrap, functools.partial):
-                unwrap = unwrap.func
-                continue
-            break
-        if hasattr(unwrap, "__globals__"):
-            obj_globals = unwrap.__globals__
-
-    if globals is None:
-        globals = obj_globals
-    if locals is None:
-        locals = obj_locals
-
-    return_value = {
-        key: value if not isinstance(value, str) else eval(value, globals, locals)
-        for key, value in ann.items()
-    }
-    return return_value
diff --git a/src/python/pants/util/typing.py b/src/python/pants/util/typing.py
new file mode 100644
--- /dev/null
+++ b/src/python/pants/util/typing.py
@@ -0,0 +1,34 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+#
+import re
+import typing
+
+ForwardRefPristine = typing.ForwardRef
+
+_union_exp = r"^([^| \[\]]*)\s*\|\s*([^\[\]]*)$"
+
+
+def _translate_piped_types_to_union(value: str) -> str:
+    # Very naive limited to top-level plain types two-legged unions only.
+    return re.sub(_union_exp, r"Union[\1, \2]", value)
+
+
+def patch_forward_ref() -> None:
+    typing.ForwardRef = ForwardRefPatched  # type: ignore[misc]
+
+
+def restore_forward_ref() -> None:
+    typing.ForwardRef = ForwardRefPristine  # type: ignore[misc]
+
+
+# We are not supposed to subclass this... but we want to support | annotations.
+class ForwardRefPatched(typing.ForwardRef, _root=True):  # type: ignore[call-arg, misc]
+    def __init__(self, arg, *args, **kwargs):
+        unionised_arg = _translate_piped_types_to_union(arg)
+        super().__init__(unionised_arg, *args, **kwargs)
+
+    def _evaluate(self, globalns, *args, **kwargs):
+        if globalns and "Union" not in globalns:
+            globalns["Union"] = typing.Union
+        return super()._evaluate(globalns, *args, **kwargs)


Test patch:
diff --git a/src/python/pants/engine/internals/rule_visitor_test.py b/src/python/pants/engine/internals/rule_visitor_test.py
--- a/src/python/pants/engine/internals/rule_visitor_test.py
+++ b/src/python/pants/engine/internals/rule_visitor_test.py
@@ -289,7 +289,7 @@ async def myrule(request: Request):
         /.*/rule_visitor_test\.py:\d+: Could not resolve type for `request\.bad_meth`
         in module pants\.engine\.internals\.rule_visitor_test\.
 
-        Return type annotation required for `bad_meth` in /.*/rule_visitor_test\.py:\d+
+        Failed to look up return type hint for `bad_meth` in /.*/rule_visitor_test\.py:\d+
         """
     )
     with pytest.raises(RuleTypeError, match=err):
@@ -304,3 +304,15 @@ async def myrule(request: Request):
         Get(str, closure_func())
 
     assert_awaitables(myrule, [(str, int)])
+
+
+class McUnion:
+    b: bool
+    v: int | float
+
+
+def test_union_types() -> None:
+    async def somerule(mc: McUnion):
+        Get(str, mc.b)
+
+    assert_awaitables(somerule, [(str, bool)])
diff --git a/src/python/pants/util/typing_test.py b/src/python/pants/util/typing_test.py
new file mode 100644
--- /dev/null
+++ b/src/python/pants/util/typing_test.py
@@ -0,0 +1,36 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from __future__ import annotations
+
+import typing
+from typing import get_type_hints
+
+import pytest
+
+from pants.util.typing import ForwardRefPatched, ForwardRefPristine, _translate_piped_types_to_union
+
+
+@pytest.mark.parametrize(
+    "value, union",
+    [
+        ("FooBar", "FooBar"),
+        ("Foo[Bar]", "Foo[Bar]"),
+        ("Foo | Bar", "Union[Foo, Bar]"),
+        ("Foo[Bar] | Foo[Baz]", "Foo[Bar] | Foo[Baz]"),
+        ("Foo[Bar | Baz]", "Foo[Bar | Baz]"),
+    ],
+)
+def test_translate_piped_types_to_union(value: str, union: str) -> None:
+    assert union == _translate_piped_types_to_union(value)
+
+
+def test_get_type_hints(monkeypatch) -> None:
+    class A:
+        b: int | float
+
+    monkeypatch.setattr(typing, "ForwardRef", ForwardRefPristine)
+    with pytest.raises(TypeError):
+        get_type_hints(A)
+
+    monkeypatch.setattr(typing, "ForwardRef", ForwardRefPatched)
+    assert get_type_hints(A)