Make stringifiers create unions if create_unions is True

python · JelleZijlstra · May 4, 2025 · Apr 22, 2025 · Apr 22, 2025 · Apr 22, 2025
commit d8109538837278622b4f3809d5ae51af942ad9ac
@@ -392,12 +392,19 @@ def binop(self, other):
    __mod__ = _make_binop(ast.Mod())
    __lshift__ = _make_binop(ast.LShift())
    __rshift__ = _make_binop(ast.RShift())
-    __or__ = _make_binop(ast.BitOr())
    __xor__ = _make_binop(ast.BitXor())
    __and__ = _make_binop(ast.BitAnd())
    __floordiv__ = _make_binop(ast.FloorDiv())
    __pow__ = _make_binop(ast.Pow())

+    def __or__(self, other):
+        if self.__stringifier_dict__.create_unions:
+            return types.UnionType[self, other]
+
+        return self.__make_new(
+            ast.BinOp(self.__get_ast(), ast.BitOr(), self.__convert_to_ast(other))
+        )
+
    del _make_binop

    def _make_rbinop(op: ast.AST):
@@ -416,12 +423,19 @@ def rbinop(self, other):
    __rmod__ = _make_rbinop(ast.Mod())
    __rlshift__ = _make_rbinop(ast.LShift())
    __rrshift__ = _make_rbinop(ast.RShift())
-    __ror__ = _make_rbinop(ast.BitOr())
    __rxor__ = _make_rbinop(ast.BitXor())
    __rand__ = _make_rbinop(ast.BitAnd())
    __rfloordiv__ = _make_rbinop(ast.FloorDiv())
    __rpow__ = _make_rbinop(ast.Pow())

+    def __ror__(self, other):
+        if self.__stringifier_dict__.create_unions:
+            return types.UnionType[other, self]
+
+        return self.__make_new(
+            ast.BinOp(self.__convert_to_ast(other), ast.BitOr(), self.__get_ast())
+        )
+
    del _make_rbinop

    def _make_compare(op):
@@ -459,12 +473,13 @@ def unary_op(self):


 class _StringifierDict(dict):
-    def __init__(self, namespace, globals=None, owner=None, is_class=False):
+    def __init__(self, namespace, globals=None, owner=None, is_class=False, create_unions=False):
        super().__init__(namespace)
        self.namespace = namespace
        self.globals = globals
        self.owner = owner
        self.is_class = is_class
+        self.create_unions = create_unions
        self.stringifiers = []

    def __missing__(self, key):
@@ -569,7 +584,13 @@ def call_annotate_function(annotate, format, *, owner=None, _is_evaluate=False):
        # that returns a bool and an defined set of attributes.
        namespace = {**annotate.__builtins__, **annotate.__globals__}
        is_class = isinstance(owner, type)
-        globals = _StringifierDict(namespace, annotate.__globals__, owner, is_class)
+        globals = _StringifierDict(
+            namespace, 
+            annotate.__globals__, 
+            owner, 
+            is_class, 
+            create_unions=True
+        )
        if annotate.__closure__:
            freevars = annotate.__code__.co_freevars
            new_closure = []