diff --git a/compiler/util/ir_data_fields.py b/compiler/util/ir_data_fields.py
index 8db01f2..fe9ef75 100644
--- a/compiler/util/ir_data_fields.py
+++ b/compiler/util/ir_data_fields.py
@@ -107,7 +107,7 @@
 
 
 class TemporaryCopyValuesList(NamedTuple):
-  """Class used to temporarily hold a CopyValueList while copying and
+  """Class used to temporarily hold a CopyValuesList while copying and
 
   constructing an IR dataclass.
   """
@@ -197,15 +197,6 @@
       if isinstance(type, v.__class__) and _is_ir_dataclass(v)
   )
 
-
-def _all_field_specs(mod):
-  specs = {
-      ir_class: FilteredIrFieldSpecs(_field_specs(ir_class))
-      for ir_class in all_ir_classes(mod)
-  }
-  return specs
-
-
 class IrDataclassSpecs:
   """Maintains a cache of all IR dataclass specs."""
 
@@ -214,8 +205,10 @@
   @classmethod
   def get_mod_specs(cls, mod):
     """Gets the IR dataclass specs for the given module."""
-    mod_specs = _all_field_specs(mod)
-    return mod_specs
+    return {
+      ir_class: FilteredIrFieldSpecs(_field_specs(ir_class))
+      for ir_class in all_ir_classes(mod)
+    }
 
   @classmethod
   def get_specs(cls, data_class):
@@ -255,17 +248,19 @@
     # Get the wrapped types if there are any
     args = get_args(type_hint)
     if origin is not None:
-      # Extract the type
+      # Extract the type.
+      type_hint = args[0]
+
+      # Underneath the hood `typing.Optional` is just a `Union[T, None]` so we
+      # have to check if it's a `Union` instead of just using `Optional`.
       if origin == Union:
-        # Extract Optionals -> Union[T, None]
+        # Make sure this is an `Optional` and not another `Union` type.
         assert len(args) == 2 and args[1] == type(None)
         container_type = FieldContainer.OPTIONAL
-        type_hint = args[0]
-        origin = get_origin(type_hint)
-
-      if origin == list:
+      elif origin == list:
         container_type = FieldContainer.LIST
-        type_hint = get_args(type_hint)[0]
+      else:
+        raise TypeError(f"Field has invalid container type: {origin}")
 
     # Resolve any forward references.
     if isinstance(type_hint, str):