Merge pull request #2468 from youngwki/typehint_model

Added return types for Model class

openmc/model/model.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from collections.abc import Iterable
 from contextlib import contextmanager
 from functools import lru_cache
@@ -7,6 +8,7 @@ from numbers import Integral
 from tempfile import NamedTemporaryFile
 import warnings
 from xml.etree import ElementTree as ET
+from typing import Optional, Dict
 
 import h5py
 
@@ -93,27 +95,27 @@ class Model:
             self.plots = plots
 
     @property
-    def geometry(self):
+    def geometry(self) -> Optional[openmc.Geometry]:
         return self._geometry
 
     @property
-    def materials(self):
+    def materials(self) -> Optional[openmc.Materials]:
         return self._materials
 
     @property
-    def settings(self):
+    def settings(self) -> Optional[openmc.Settings]:
         return self._settings
 
     @property
-    def tallies(self):
+    def tallies(self) -> Optional[openmc.Tallies]:
         return self._tallies
 
     @property
-    def plots(self):
+    def plots(self) -> Optional[openmc.Plots]:
         return self._plots
 
     @property
-    def is_initialized(self):
+    def is_initialized(self) -> bool:
         try:
             import openmc.lib
             return openmc.lib.is_initialized
@@ -122,7 +124,7 @@ class Model:
 
     @property
     @lru_cache(maxsize=None)
-    def _materials_by_id(self):
+    def _materials_by_id(self) -> dict:
         """Dictionary mapping material ID --> material"""
         if self.materials:
             mats = self.materials
@@ -132,14 +134,14 @@ class Model:
 
     @property
     @lru_cache(maxsize=None)
-    def _cells_by_id(self):
+    def _cells_by_id(self) -> dict:
         """Dictionary mapping cell ID --> cell"""
         cells = self.geometry.get_all_cells()
         return {cell.id: cell for cell in cells.values()}
 
     @property
     @lru_cache(maxsize=None)
-    def _cells_by_name(self):
+    def _cells_by_name(self) -> Dict[int, openmc.Cell]:
         # Get the names maps, but since names are not unique, store a set for
         # each name key. In this way when the user requests a change by a name,
         # the change will be applied to all of the same name.
@@ -152,7 +154,7 @@ class Model:
 
     @property
     @lru_cache(maxsize=None)
-    def _materials_by_name(self):
+    def _materials_by_name(self) -> Dict[int, openmc.Material]:
         if self.materials is None:
             mats = self.geometry.get_all_materials().values()
         else:
@@ -207,7 +209,7 @@ class Model:
     @classmethod
     def from_xml(cls, geometry='geometry.xml', materials='materials.xml',
                  settings='settings.xml', tallies='tallies.xml',
-                 plots='plots.xml'):
+                 plots='plots.xml') -> Model:
         """Create model from existing XML files
 
         Parameters
@@ -604,7 +606,6 @@ class Model:
         this method creates the XML files and runs OpenMC via a system call. In
         both cases this method returns the path to the last statepoint file
         generated.
-
         .. versionchanged:: 0.12
             Instead of returning the final k-effective value, this function now
             returns the path to the final statepoint written.