From fc00c81f28ac506b3ae72590f428294027cd2a18 Mon Sep 17 00:00:00 2001
From: Matej Aleksandrov <maleksandrov@google.com>
Date: Fri, 19 Jun 2026 13:05:36 -0700
Subject: [PATCH] Automated Code Change

PiperOrigin-RevId: 935024351
---
 fiddle/_src/codegen/auto_config/ir_printer.py | 27 ++++++++++++++++---
 .../codegen/auto_config/ir_printer_test.py    |  2 +-
 fiddle/_src/validation/check_types_test.py    |  8 +++---
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/fiddle/_src/codegen/auto_config/ir_printer.py b/fiddle/_src/codegen/auto_config/ir_printer.py
index f81c8819..5c732f3f 100644
--- a/fiddle/_src/codegen/auto_config/ir_printer.py
+++ b/fiddle/_src/codegen/auto_config/ir_printer.py
@@ -20,6 +20,7 @@
 """
 
 import inspect
+import types
 import typing
 from typing import Any, List
 
@@ -28,9 +29,27 @@
 from fiddle._src.codegen.auto_config import code_ir
 import libcst as cst
 
-# Values from the `typing` module. This probably pulls in too much stuff, but
-# we don't have to be 100% precise in this module, it's just for debug printing.
-typing_consts = [getattr(typing, name) for name in dir(typing)]
+
+def _is_typing_or_generic(value: Any) -> bool:
+  """Checks if a value is a typing construct or generic alias."""
+  union_type = getattr(types, "UnionType", None)
+  generic_alias = getattr(types, "GenericAlias", None)
+
+  types_to_check = []
+  if union_type is not None:
+    types_to_check.append(union_type)
+  if generic_alias is not None:
+    types_to_check.append(generic_alias)
+
+  if types_to_check and isinstance(value, tuple(types_to_check)):
+    return True
+
+  if hasattr(value, "__module__") and value.__module__ == "typing":
+    return True
+  t = type(value)
+  if hasattr(t, "__module__") and t.__module__ == "typing":
+    return True
+  return False
 
 
 def format_py_reference(value: Any) -> str:
@@ -121,7 +140,7 @@ def traverse(value, state: daglish.State) -> str:
       return value.value
     elif isinstance(value, type) and value is not typing.Any:
       return value.__name__
-    elif value in typing_consts or type(value) in typing_consts:
+    elif _is_typing_or_generic(value):
       return str(value)
     else:
       return f"<<<custom:{repr(value)}>>>"
diff --git a/fiddle/_src/codegen/auto_config/ir_printer_test.py b/fiddle/_src/codegen/auto_config/ir_printer_test.py
index 095f2ed6..77c7bfcc 100644
--- a/fiddle/_src/codegen/auto_config/ir_printer_test.py
+++ b/fiddle/_src/codegen/auto_config/ir_printer_test.py
@@ -45,7 +45,7 @@ def test_format_expr_primitives(self):
     self.assertIn(
         ir_printer.format_expr(Optional[str]),
         # This differs based on the Python verison.
-        {"typing.Union[str, NoneType]", "typing.Optional[str]"},
+        {"typing.Union[str, NoneType]", "typing.Optional[str]", "str | None"},
     )
 
   def test_format_containers(self):
diff --git a/fiddle/_src/validation/check_types_test.py b/fiddle/_src/validation/check_types_test.py
index d935daa4..54d60eb3 100644
--- a/fiddle/_src/validation/check_types_test.py
+++ b/fiddle/_src/validation/check_types_test.py
@@ -201,10 +201,10 @@ def test_union_and_optional_types(self):
     )
     with self.assertRaisesRegex(
         TypeError,
-        ".*For attribute .*union_field provided type:"
-        " .*BadTokenEmbedder.* is not of annotated/declared"
-        " type: typing.Union.*str,"
-        " fiddle._src.testing.example.fake_encoder_decoder.TokenEmbedder.*",
+        r".*For attribute .*union_field provided type:"
+        r" .*BadTokenEmbedder.* is not of annotated/declared"
+        r" type: (typing.Union\[str,\s*|str\s*\|\s*)"
+        r"fiddle\._src\.testing\.example\.fake_encoder_decoder\.TokenEmbedder.*",
     ):
       check_types.check_types(cfg)