From 2fa13a1a8dc687fad45e474305f381a74fa0df37 Mon Sep 17 00:00:00 2001
From: BitterB0NG0 <93603537+BitterB0NG0@users.noreply.github.com>
Date: Thu, 21 Nov 2024 00:52:45 -0700
Subject: [PATCH] utils/types: update distribution number & comments

---
 src/pymcnp/files/inp/_factory.py |  54 ++++++-
 src/pymcnp/files/inp/data.py     |   7 +-
 src/pymcnp/files/utils/types.py  | 235 ++++++++++++++++++++++---------
 3 files changed, 227 insertions(+), 69 deletions(-)

diff --git a/src/pymcnp/files/inp/_factory.py b/src/pymcnp/files/inp/_factory.py
index cbede80..a197355 100644
--- a/src/pymcnp/files/inp/_factory.py
+++ b/src/pymcnp/files/inp/_factory.py
@@ -231,6 +231,33 @@ def build(self):
 
                 o += '\n'
 
+            elif attribute.type.startswith('Union'):
+                inside_types = attribute.type[len('Union[') : -1].split(', ')
+
+                a = True
+                for inside_type in inside_types:
+                    o += (
+                        f'        {'if' if a else 'elif'} isinstance(inside_type, {inside_type}):\n'
+                    )
+
+                    if inside_type.startswith('tuple'):
+                        inside_inside_type = inside_type[len('tuple[') : -1]
+
+                        if inside_inside_type == 'str':
+                            # Union[tuple[str], ...]
+                            o += f'            {attribute.name} = [tokens.popl() for _ in range(0, len(tokens))]\n'
+                        else:
+                            # Union[tuple[?], ...]
+                            o += f'            {attribute.name} = [{inside_inside_type}.popl() for _ in range(0, len(tokens))]\n'
+                    elif inside_type == 'str':
+                        # Union[str, ...]
+                        o += f'            {attribute.name} = tokens.popl()\n'
+                    else:
+                        # Union[?, ...]
+                        o += f'            {attribute.name} = {inside_type}.from_mcnp(tokens.popl())\n'
+
+                    a = False
+
             elif attribute.type.startswith('tuple'):
                 inside_type = attribute.type[len('tuple[') : -1]
 
@@ -256,7 +283,7 @@ def build(self):
                 entry_attributes = len(self.entries[index].attributes)
                 o += f'        {attribute.name} = {attribute.type}.from_mcnp(" ".join([tokens.popl() for _ in range(0, {entry_attributes})]))\n'
             else:
-                # object
+                # ?
                 o += f'        {attribute.name} = {attribute.type}.from_mcnp(tokens.popl())\n'
 
         o += '\n'
@@ -598,6 +625,31 @@ def build(self, name, mnemonic):
         elif self.attribute.type == 'str':
             o += '        value = tokens.popl()\n'
 
+        elif self.attribute.type.startswith('Union'):
+            inside_types = self.attribute.type[len('Union[') : -1].split(', ')
+
+            a = True
+            for inside_type in inside_types:
+                o += f'        {'if' if a else 'elif'} isinstance(inside_type, {inside_type}):\n'
+
+                if inside_type.startswith('tuple'):
+                    inside_inside_type = inside_type[len('tuple[') : -1]
+
+                    if inside_inside_type == 'str':
+                        # Union[tuple[str], ...]
+                        o += f'            {self.attribute.name} = [tokens.popl() for _ in range(0, len(tokens))]\n'
+                    else:
+                        # Union[tuple[?], ...]
+                        o += f'            {self.attribute.name} = [{inside_inside_type}.popl() for _ in range(0, len(tokens))]\n'
+                elif inside_type == 'str':
+                    # Union[str, ...]
+                    o += f'            {self.attribute.name} = tokens.popl()\n'
+                else:
+                    # Union[?, ...]
+                    o += f'            {self.attribute.name} = {inside_type}.from_mcnp(tokens.popl())\n'
+
+                a = False
+
         else:
             o += f'        value = {self.attribute.type}.from_mcnp(tokens.popl())\n'
 
diff --git a/src/pymcnp/files/inp/data.py b/src/pymcnp/files/inp/data.py
index 275cf7e..9651483 100644
--- a/src/pymcnp/files/inp/data.py
+++ b/src/pymcnp/files/inp/data.py
@@ -1705,7 +1705,12 @@ def to_mcnp(self) -> str:
         ),
         _factory.DataOptionFactory(
             'tme',
-            _factory.AttributeFactory('time', 'types.McnpReal', 'Time in shakes', 'time >= 0'),
+            _factory.AttributeFactory(
+                'time',
+                'Union[types.McnpReal, types.EmbeddedDistributionNumber]',
+                'Time in shakes',
+                'time >= 0',
+            ),
         ),
         _factory.DataOptionFactory(
             'dir',
diff --git a/src/pymcnp/files/utils/types.py b/src/pymcnp/files/utils/types.py
index c3519d8..21fbb1f 100644
--- a/src/pymcnp/files/utils/types.py
+++ b/src/pymcnp/files/utils/types.py
@@ -1,9 +1,5 @@
 """
-``types`` contains class representing basic MCNP types.
-
-``types`` packages the ``McnpInteger``, ``McnpReal``, ``Zaid``, and
-``Designator`` classes, providing an object-oriented, importable interface for
-MCNP types.
+Contains classes representing basic MCNP value types.
 """
 
 from __future__ import annotations
@@ -11,13 +7,16 @@
 import enum
 from typing import Literal, Final
 
-from . import _parser
 from . import errors
+from . import _parser
+from . import _object
 
 
-class DistributionNumber:
+class DistributionNumber(_object.PyMcnpObject):
     """
-    ``DistributionNumber`` represents MCNP distribution numbers.
+    Represents MCNP distribution numbers.
+
+    ``DistributionNumber`` implements ``_object.PyMcnpObject``.
 
     Attributes:
         n: number.
@@ -25,7 +24,7 @@ class DistributionNumber:
 
     def __init__(self, n: int):
         """
-        ``__init__`` initializes ``DistributionNumber``.
+        Initializes ``DistributionNumber``.
 
         Parameters:
             n: number.
@@ -42,14 +41,12 @@ def __init__(self, n: int):
     @staticmethod
     def from_mcnp(source: str):
         """
-        ``from_mcnp`` generates ``DistributionNumber`` objects from INP.
+        Generates ``DistributionNumber`` objects from INP.
 
-        ``from_mcnp`` constructs instances of ``DistributionNumber`` from INP
-        source strings, so it operates as a class constructor method
-        and INP parser helper function.
+        ``from_mcnp`` translates from INP to PyMCNP; it parses INP.
 
         Parameters:
-            source: INP for distribution number.
+            source: INP for ``DistributionNumber``.
 
         Returns:
             ``DistributionNumber`` object.
@@ -66,10 +63,93 @@ def from_mcnp(source: str):
 
         return DistributionNumber(int(match[1]))
 
+    def to_mcnp(self):
+        """
+        Generates INP from ``DistributionNumber`` objects.
+
+        ``to_mcnp`` translates from PyMCNP to INP.
+
+        Returns:
+            INP for ``DistributionNumber``.
+        """
+
+        return f'd{self.n}'
 
-class Zaid:
+
+class EmbeddedDistributionNumber(_object.PyMcnpObject):
     """
-    ``Zaid`` represents nuclide information numbers.
+    Represents MCNP embedded distribution numbers.
+
+    ``EmbeddedDistributionNumber`` implements ``_object.PyMcnpObject``.
+
+    Attributes:
+        numbers: Tuple of distribution numbers.
+    """
+
+    def __init__(self, numbers: tuple[DistributionNumber]):
+        """
+        Initializes ``EmbeddedDistributionNumber``.
+
+        Parameters:
+            numbers: Tuple of distribution numbers.
+
+        Raises:
+            McnpError: INVALID_DN.
+        """
+
+        if numbers is None:
+            raise errors.McnpError(errors.McnpCode.INVALID_DN, str(numbers))
+
+        for number in numbers:
+            if number is None:
+                raise errors.McnpError(errors.McnpCode.INVALID_DN, str(number))
+
+        self.numbers: Final[tuple[DistributionNumber]] = numbers
+
+    @staticmethod
+    def from_mcnp(source: str):
+        """
+        Generates ``EmbeddedDistributionNumber`` objects from INP.
+
+        ``from_mcnp`` translates from INP to PyMCNP; it parses INP.
+
+        Parameters:
+            source: INP for ``EmbeddedDistributionNumber``.
+
+        Returns:
+            ``EmbeddedDistributionNumber`` object.
+
+        Raises:
+            McnpError: UNRECOGNIZED_KEYWORD.
+        """
+
+        source = _parser.Preprocessor.process_inp(source)
+        tokens = re.split(r'>', source)
+
+        numbers = []
+        for token in tokens:
+            numbers.append(DistributionNumber.from_mcnp(token))
+
+        return EmbeddedDistributionNumber(tuple(numbers))
+
+    def to_mcnp(self):
+        """
+        Generates INP from ``EmbeddedDistributionNumber`` objects.
+
+        ``to_mcnp`` translates from PyMCNP to INP.
+
+        Returns:
+            INP for ``EmbeddedDistributionNumber``.
+        """
+
+        return '>'.join(number.to_mcnp() for number in self.numbers)
+
+
+class Zaid(_object.PyMcnpObject):
+    """
+    Represents nuclide information numbers.
+
+    ``Zaid`` implements ``_object.PyMcnpObject``.
 
     Attributes:
         z: Atomic number.
@@ -79,7 +159,7 @@ class Zaid:
 
     def __init__(self, z: int, a: int, abx: str = None):
         """
-        ``__init__`` initializes ``Zaid``.
+        Initializes ``Zaid``.
 
         Parameters:
             z: Atomic number.
@@ -103,14 +183,12 @@ def __init__(self, z: int, a: int, abx: str = None):
     @staticmethod
     def from_mcnp(source: str):
         """
-        ``from_mcnp`` generates ``Zaid`` objects from INP.
+        Generates ``Zaid`` objects from INP.
 
-        ``from_mcnp`` constructs instances of ``Zaid`` from INP
-        source strings, so it operates as a class constructor method
-        and INP parser helper function.
+        ``from_mcnp`` translates from INP to PyMCNP; it parses INP.
 
         Parameters:
-            source: INP for zaid.
+            source: INP for ``Zaid``.
 
         Returns:
             ``Zaid`` object.
@@ -147,13 +225,12 @@ def from_mcnp(source: str):
 
     def to_mcnp(self) -> str:
         """
-        ``to_mcnp`` generates INP from ``Zaid`` objects.
+        Generates INP from ``Zaid`` objects.
 
-        ``to_mcnp`` creates INP source string from ``Zaid``
-        objects, so it provides an MCNP endpoint.
+        ``to_mcnp`` translates from PyMCNP to INP.
 
         Returns:
-            INP string for ``Zaid`` object.
+            INP for ``Zaid``.
         """
 
         return (
@@ -166,9 +243,11 @@ def __str__(self):
         return self.to_mcnp()
 
 
-class Particle(str, enum.Enum):
+class Particle(_object.PyMcnpObject, enum.StrEnum):
     """
-    ``Particle`` represents individular particle designators.
+    Represents particle designators.
+
+    ``Particle`` implements ``_object.PyMcnpObject`` and ``enum.StrEnum``.
     """
 
     NEUTRON = 'n'
@@ -209,10 +288,40 @@ class Particle(str, enum.Enum):
     ALPHA = 'a'
     HEAVY_IONS = '#'
 
+    @staticmethod
+    def from_mcnp(source: str):
+        """
+        Generates ``Particle`` objects from INP.
+
+        ``from_mcnp`` translates from INP to PyMCNP; it parses INP.
 
-class Designator:
+        Parameters:
+            source: INP for ``Particle``.
+
+        Returns:
+            ``Particle`` object.
+        """
+
+        return Particle(source)
+
+    def to_mcnp(self):
+        """
+        Generates INP from ``Particle`` objects.
+
+        ``to_mcnp`` translates from PyMCNP to INP.
+
+        Returns:
+            INP for ``Particle``.
+        """
+
+        return self.value
+
+
+class Designator(_object.PyMcnpObject):
     """
-    ``Designator`` represents MCNP particle designators.
+    Represents MCNP particle designators.
+
+    ``Designator`` implements ``_object.PyMcnpObject``.
 
     Attributes:
         particles: Tuple of particles.
@@ -220,7 +329,7 @@ class Designator:
 
     def __init__(self, particles: tuple[Particle]):
         """
-        ``__init__`` initializes ``Designator``.
+        Initializes ``Designator``.
 
         Parameters:
             particles: Tuple of particles.
@@ -244,17 +353,15 @@ def __init__(self, particles: tuple[Particle]):
     @staticmethod
     def from_mcnp(source: str):
         """
-        ``from_mcnp`` generates ``Designator`` objects from INP.
+        Generates ``Designator`` objects from INP.
 
-        ``from_mcnp`` constructs instances of ``Designator`` from INP
-        source strings, so it operates as a class constructor method
-        and INP parser helper function.
+        ``from_mcnp`` translates from INP to PyMCNP; it parses INP.
 
         Parameters:
-            source: INP for particle designator(s).
+            source: INP for ``Designator``.
 
         Returns:
-            Tuple of ``Designator`` objects.
+            ``Designator`` object.
         """
 
         try:
@@ -266,24 +373,22 @@ def from_mcnp(source: str):
 
     def to_mcnp(self) -> str:
         """
-        ``to_mcnp`` generates INP from ``Designator`` objects.
+        Generates INP from ``Designator`` objects.
 
-        ``to_mcnp`` creates INP source string from ``Designator``
-        objects, so it provides an MCNP endpoint.
+        ``to_mcnp`` translates from PyMCNP to INP.
 
         Returns:
-            INP string for ``Designator`` object.
+            INP for ``Designator``.
         """
 
         return ','.join(Particle(particle) for particle in self.particles)
 
-    def __eq__(self, other):
-        return self.particles == other.particles
-
 
-class McnpInteger:
+class McnpInteger(_object.PyMcnpObject):
     """
-    ``McnpInteger`` represents the INP integer value.
+    Represents the INP integer value.
+
+    ``McnpInteger`` implements ``_object.PyMcnpObject``.
 
     Attributes:
         value: Integer or J jump symbol.
@@ -291,7 +396,7 @@ class McnpInteger:
 
     def __init__(self, value: int | Literal['j']):
         """
-        ``__init__`` initializes ``McnpInteger``.
+        Initializes ``McnpInteger``.
 
         Parameters:
             value: Integer or J jump symbol.
@@ -315,11 +420,9 @@ def __init__(self, value: int | Literal['j']):
     @staticmethod
     def from_mcnp(source: str):
         """
-        ``from_mcnp`` generates ``McnpInteger`` objects from INP.
+        Generates ``McnpInteger`` objects from INP.
 
-        ``from_mcnp`` constructs instances of ``McnpInteger`` from INP
-        source strings, so it operates as a class constructor method
-        and INP parser helper function.
+        ``from_mcnp`` translates from INP to PyMCNP.
 
         Parameters:
             source: INP for value.
@@ -349,13 +452,12 @@ def from_mcnp(source: str):
 
     def to_mcnp(self):
         """
-        ``to_mcnp`` generates INP from ``McnpInteger`` objects.
+        Generates INP from ``McnpInteger`` objects.
 
-        ``to_mcnp`` creates INP source string from ``McnpInteger``
-        objects, so it provides an MCNP endpoint.
+        ``to_mcnp`` translates from PyMCNP to INP.
 
         Returns:
-            INP string for ``McnpInteger`` object.
+            INP for ``McnpInteger``.
         """
 
         return str(self.value)
@@ -409,9 +511,11 @@ def __int__(self) -> int:
         return int(self.value)
 
 
-class McnpReal:
+class McnpReal(_object.PyMcnpObject):
     """
-    ``McnpReal`` represents the INP real/floating-point value.
+    Represents the INP real/floating-point value.
+
+    ``McnpReal`` implements ``_object.PyMcnpObject``.
 
     Attributes:
         value: Floating-point number or J jump symbol.
@@ -421,7 +525,7 @@ class McnpReal:
 
     def __init__(self, value: float | Literal['j']):
         """
-        ``__init__`` initializes ``McnpReal``.
+        Initializes ``McnpReal``.
 
         Parameters:
             value: Floating-point number or J jump symbol.
@@ -445,14 +549,12 @@ def __init__(self, value: float | Literal['j']):
     @staticmethod
     def from_mcnp(source: str):
         """
-        ``from_mcnp`` generates ``McnpReal`` objects from INP.
+        Generates ``McnpReal`` objects from INP.
 
-        ``from_mcnp`` constructs instances of ``McnpReal`` from INP
-        source strings, so it operates as a class constructor method
-        and INP parser helper function.
+        ``from_mcnp`` translates from INP to PyMCNP; it parses INP.
 
         Parameters:
-            source: INP for value.
+            source: INP for ``McnpReal``.
 
         Returns:
             ``McnpReal`` object.
@@ -480,13 +582,12 @@ def from_mcnp(source: str):
 
     def to_mcnp(self):
         """
-        ``to_mcnp`` generates INP from ``McnpReal`` objects.
+        Generates INP from ``McnpReal`` objects.
 
-        ``to_mcnp`` creates INP source string from ``McnpReal``
-        objects, so it provides an MCNP endpoint.
+        ``to_mcnp`` translates from PyMCNP to INP.
 
         Returns:
-            INP string for ``McnpReal`` object.
+            INP for ``McnpReal``.
         """
 
         return str(self.value)