Merge pull request #2273 from shimwell/type_hinting_for_source.py

Type hinting for source.py

openmc/__init__.py

@@ -10,11 +10,11 @@ from openmc.material import *
 from openmc.plots import *
 from openmc.region import *
 from openmc.volume import *
-from openmc.source import *
 from openmc.weight_windows import *
-from openmc.settings import *
 from openmc.surface import *
 from openmc.universe import *
+from openmc.source import *
+from openmc.settings import *
 from openmc.lattice import *
 from openmc.filter import *
 from openmc.filter_expansion import *

openmc/checkvalue.py

@@ -1,12 +1,12 @@
 import copy
 import os
-from typing import Union
+import typing  # required to prevent typing.Union namespace overwriting Union
 from collections.abc import Iterable
 
 import numpy as np
 
 # Type for arguments that accept file paths
-PathLike = Union[str, os.PathLike]
+PathLike = typing.Union[str, os.PathLike]
 
 
 def check_type(name, value, expected_type, expected_iter_type=None, *, none_ok=False):

openmc/deplete/results.py

@@ -1,7 +1,8 @@
 import numbers
 import bisect
 import math
-from typing import Iterable, Optional, Tuple, Union
+import typing  # required to prevent typing.Union namespace overwriting Union
+from typing import Iterable, Optional, Tuple
 from warnings import warn
 
 import h5py
@@ -95,7 +96,7 @@ class Results(list):
 
     def get_atoms(
         self,
-        mat: Union[Material, str],
+        mat: typing.Union[Material, str],
         nuc: str,
         nuc_units: str = "atoms",
         time_units: str = "s"
@@ -165,7 +166,7 @@ class Results(list):
 
     def get_reaction_rate(
         self,
-        mat: Union[Material, str],
+        mat: typing.Union[Material, str],
         nuc: str,
         rx: str
     ) -> Tuple[np.ndarray, np.ndarray]:

openmc/material.py

@@ -6,7 +6,7 @@ from pathlib import Path
 import re
 import typing  # imported separately as py3.8 requires typing.Iterable
 import warnings
-from typing import Optional, Union
+from typing import Optional
 from xml.etree import ElementTree as ET
 
 import h5py
@@ -968,7 +968,7 @@ class Material(IDManagerMixin):
 
         Returns
         -------
-        Union[dict, float]
+        typing.Union[dict, float]
             If by_nuclide is True then a dictionary whose keys are nuclide
             names and values are activity is returned. Otherwise the activity
             of the material is returned as a float.

openmc/settings.py

@@ -6,7 +6,8 @@ import itertools
 from math import ceil
 from numbers import Integral, Real
 from pathlib import Path
-from typing import Optional, Union
+import typing  # required to prevent typing.Union namespace overwriting Union
+from typing import Optional
 from xml.etree import ElementTree as ET
 
 import openmc.checkvalue as cv
@@ -582,7 +583,7 @@ class Settings:
         self._max_order = max_order
 
     @source.setter
-    def source(self, source: Union[Source, typing.Iterable[Source]]):
+    def source(self, source: typing.Union[Source, typing.Iterable[Source]]):
         if not isinstance(source, MutableSequence):
             source = [source]
         self._source = cv.CheckedList(Source, 'source distributions', source)
@@ -843,7 +844,7 @@ class Settings:
 
     @volume_calculations.setter
     def volume_calculations(
-        self, vol_calcs: Union[VolumeCalculation, typing.Iterable[VolumeCalculation]]
+        self, vol_calcs: typing.Union[VolumeCalculation, typing.Iterable[VolumeCalculation]]
     ):
         if not isinstance(vol_calcs, MutableSequence):
             vol_calcs = [vol_calcs]
@@ -889,7 +890,7 @@ class Settings:
         self._write_initial_source = value
 
     @weight_windows.setter
-    def weight_windows(self, value: Union[WeightWindows, typing.Iterable[WeightWindows]]):
+    def weight_windows(self, value: typing.Union[WeightWindows, typing.Iterable[WeightWindows]]):
         if not isinstance(value, MutableSequence):
             value = [value]
         self._weight_windows = cv.CheckedList(WeightWindows, 'weight windows', value)

openmc/source.py

@@ -2,6 +2,9 @@ from collections.abc import Iterable
 from enum import Enum
 from numbers import Real
 import warnings
+import typing  # imported separately as py3.8 requires typing.Iterable
+# also required to prevent typing.Union namespace overwriting Union
+from typing import Optional, Sequence
 from xml.etree import ElementTree as ET
 
 import numpy as np
@@ -9,6 +12,7 @@ import h5py
 
 import openmc
 import openmc.checkvalue as cv
+from openmc.checkvalue import PathLike
 from openmc.stats.multivariate import UnitSphere, Spatial
 from openmc.stats.univariate import Univariate
 from ._xml import get_text
@@ -70,9 +74,19 @@ class Source:
 
     """
 
-    def __init__(self, space=None, angle=None, energy=None, time=None, filename=None,
-                 library=None, parameters=None, strength=1.0, particle='neutron',
-                 domains=None):
+    def __init__(
+        self,
+        space: Optional[openmc.stats.Spatial] = None,
+        angle: Optional[openmc.stats.UnitSphere] = None,
+        energy: Optional[openmc.stats.Univariate] = None,
+        time: Optional[openmc.stats.Univariate] = None,
+        filename: Optional[str] = None,
+        library: Optional[str] = None,
+        parameters: Optional[str] = None,
+        strength: float = 1.0,
+        particle: str = 'neutron',
+        domains: Optional[Sequence[typing.Union[openmc.Cell, openmc.Material, openmc.Universe]]] = None
+    ):
         self._space = None
         self._angle = None
         self._energy = None
@@ -209,7 +223,7 @@ class Source:
         cv.check_value('source particle', particle, ['neutron', 'photon'])
         self._particle = particle
 
-    def to_xml_element(self):
+    def to_xml_element(self) -> ET.Element:
         """Return XML representation of the source
 
         Returns
@@ -244,7 +258,7 @@ class Source:
         return element
 
     @classmethod
-    def from_xml_element(cls, elem):
+    def from_xml_element(cls, elem: ET.Element) -> 'openmc.Source':
         """Generate source from an XML element
 
         Parameters
@@ -349,8 +363,18 @@ class SourceParticle:
         Type of the particle
 
     """
-    def __init__(self, r=(0., 0., 0.), u=(0., 0., 1.), E=1.0e6, time=0.0, wgt=1.0,
-                 delayed_group=0, surf_id=0, particle=ParticleType.NEUTRON):
+    def __init__(
+        self,
+        r: typing.Iterable[float] = (0., 0., 0.),
+        u: typing.Iterable[float] = (0., 0., 1.),
+        E: float = 1.0e6,
+        time: float = 0.0,
+        wgt: float = 1.0,
+        delayed_group: int = 0,
+        surf_id: int = 0,
+        particle: ParticleType = ParticleType.NEUTRON
+    ):
+
         self.r = tuple(r)
         self.u = tuple(u)
         self.E = float(E)
@@ -364,7 +388,7 @@ class SourceParticle:
         name = self.particle.name.lower()
         return f'<SourceParticle: {name} at E={self.E:.6e} eV>'
 
-    def to_tuple(self):
+    def to_tuple(self) -> tuple:
         """Return source particle attributes as a tuple
 
         Returns
@@ -377,7 +401,10 @@ class SourceParticle:
                 self.delayed_group, self.surf_id, self.particle.value)
 
 
-def write_source_file(source_particles, filename, **kwargs):
+def write_source_file(
+    source_particles: typing.Iterable[SourceParticle],
+    filename: PathLike, **kwargs
+):
     """Write a source file using a collection of source particles
 
     Parameters