diff --git a/pdl-compiler/tests/python_generator_test.py b/pdl-compiler/tests/python_generator_test.py
index 97e4968..21ab14b 100644
--- a/pdl-compiler/tests/python_generator_test.py
+++ b/pdl-compiler/tests/python_generator_test.py
@@ -21,7 +21,6 @@
 import enum
 import json
 import typing
-import typing_extensions
 import unittest
 from importlib import resources
 
@@ -57,12 +56,12 @@ def create_object(typ, value):
             field_type = field_types[f]
             values[f] = create_object(field_type, v)
         return typ(**values)
-    elif typing_extensions.get_origin(typ) is list:
-        typ = typing_extensions.get_args(typ)[0]
+    elif typing.get_origin(typ) is list:
+        typ = typing.get_args(typ)[0]
         return [create_object(typ, v) for v in value]
-    elif typing_extensions.get_origin(typ) is typing.Union:
+    elif typing.get_origin(typ) is typing.Union:
         # typing.Optional[int] expands to typing.Union[int, None]
-        typ = typing_extensions.get_args(typ)[0]
+        typ = typing.get_args(typ)[0]
         return create_object(typ, value) if value is not None else None
     elif typ is bytes:
         return bytes(value)