From 3692e7331fc3315b208b9e794f4ef534b7fd4664 Mon Sep 17 00:00:00 2001 From: Christina Cai Date: Wed, 12 Apr 2023 14:41:14 -0400 Subject: [PATCH] added return types for geometry.py --- openmc/geometry.py | 57 ++++++++++++++++++++++++---------------------- 1 file changed, 30 insertions(+), 27 deletions(-) diff --git a/openmc/geometry.py b/openmc/geometry.py index ad86cc56b4..88871c5441 100644 --- a/openmc/geometry.py +++ b/openmc/geometry.py @@ -1,12 +1,15 @@ +from __future__ import annotations import os -import typing +import typing from collections import OrderedDict, defaultdict from collections.abc import Iterable from copy import deepcopy from pathlib import Path -from xml.etree import ElementTree as ET import warnings +from xml.etree import ElementTree as ET +import numpy as np + import openmc import openmc._xml as xml from .checkvalue import check_type, check_less_than, check_greater_than, PathLike @@ -51,19 +54,19 @@ class Geometry: self._root_universe = univ @property - def root_universe(self): + def root_universe(self) -> openmc.UniverseBase: return self._root_universe @property - def bounding_box(self): + def bounding_box(self) -> np.ndarray: return self.root_universe.bounding_box @property - def merge_surfaces(self): + def merge_surfaces(self) -> bool: return self._merge_surfaces @property - def surface_precision(self): + def surface_precision(self) -> int: return self._surface_precision @root_universe.setter @@ -105,7 +108,7 @@ class Geometry: if universe.id in volume_calc.volumes: universe.add_volume_information(volume_calc) - def to_xml_element(self, remove_surfs=False): + def to_xml_element(self, remove_surfs=False) -> ET.Element: """Creates a 'geometry' element to be written to an XML file. Parameters @@ -164,7 +167,7 @@ class Geometry: tree.write(str(p), xml_declaration=True, encoding='utf-8') @classmethod - def from_xml_element(cls, elem, materials=None): + def from_xml_element(cls, elem, materials=None) -> openmc.Geometry: """Generate geometry from an XML element Parameters @@ -260,7 +263,7 @@ class Geometry: cls, path: PathLike = 'geometry.xml', materials: typing.Optional[typing.Union[PathLike, 'openmc.Materials']] = 'materials.xml' - ): + ) -> openmc.Geometry: """Generate geometry from XML file Parameters @@ -290,7 +293,7 @@ class Geometry: return cls.from_xml_element(root, materials) - def find(self, point): + def find(self, point) -> list: """Find cells/universes/lattices which contain a given point Parameters @@ -307,7 +310,7 @@ class Geometry: """ return self.root_universe.find(point) - def get_instances(self, paths): + def get_instances(self, paths) -> typing.Union[int, typing.List[int]]: """Return the instance number(s) for a cell/material in a geometry path. The instance numbers are used as indices into distributed @@ -354,7 +357,7 @@ class Geometry: return indices if return_list else indices[0] - def get_all_cells(self): + def get_all_cells(self) -> OrderedDict: """Return all cells in the geometry. Returns @@ -368,7 +371,7 @@ class Geometry: else: return OrderedDict() - def get_all_universes(self): + def get_all_universes(self) -> OrderedDict: """Return all universes in the geometry. Returns @@ -383,7 +386,7 @@ class Geometry: universes.update(self.root_universe.get_all_universes()) return universes - def get_all_materials(self): + def get_all_materials(self) -> OrderedDict: """Return all materials within the geometry. Returns @@ -398,7 +401,7 @@ class Geometry: else: return OrderedDict() - def get_all_material_cells(self): + def get_all_material_cells(self) -> OrderedDict: """Return all cells filled by a material Returns @@ -417,7 +420,7 @@ class Geometry: return material_cells - def get_all_material_universes(self): + def get_all_material_universes(self) -> OrderedDict: """Return all universes having at least one material-filled cell. This method can be used to find universes that have at least one cell @@ -440,7 +443,7 @@ class Geometry: return material_universes - def get_all_lattices(self): + def get_all_lattices(self) -> OrderedDict: """Return all lattices defined Returns @@ -458,7 +461,7 @@ class Geometry: return lattices - def get_all_surfaces(self): + def get_all_surfaces(self) -> OrderedDict: """ Return all surfaces used in the geometry @@ -475,7 +478,7 @@ class Geometry: surfaces = cell.region.get_surfaces(surfaces) return surfaces - def _get_domains_by_name(self, name, case_sensitive, matching, domain_type): + def _get_domains_by_name(self, name, case_sensitive, matching, domain_type) -> list: if not case_sensitive: name = name.lower() @@ -492,7 +495,7 @@ class Geometry: domains.sort(key=lambda x: x.id) return domains - def get_materials_by_name(self, name, case_sensitive=False, matching=False): + def get_materials_by_name(self, name, case_sensitive=False, matching=False) -> typing.List[openmc.Material]: """Return a list of materials with matching names. Parameters @@ -513,7 +516,7 @@ class Geometry: """ return self._get_domains_by_name(name, case_sensitive, matching, 'material') - def get_cells_by_name(self, name, case_sensitive=False, matching=False): + def get_cells_by_name(self, name, case_sensitive=False, matching=False) -> typing.List[openmc.Cell]: """Return a list of cells with matching names. Parameters @@ -534,7 +537,7 @@ class Geometry: """ return self._get_domains_by_name(name, case_sensitive, matching, 'cell') - def get_surfaces_by_name(self, name, case_sensitive=False, matching=False): + def get_surfaces_by_name(self, name, case_sensitive=False, matching=False) -> typing.List[openmc.Surface]: """Return a list of surfaces with matching names. .. versionadded:: 0.13.3 @@ -557,7 +560,7 @@ class Geometry: """ return self._get_domains_by_name(name, case_sensitive, matching, 'surface') - def get_cells_by_fill_name(self, name, case_sensitive=False, matching=False): + def get_cells_by_fill_name(self, name, case_sensitive=False, matching=False) -> typing.List[openmc.Cell]: """Return a list of cells with fills with matching names. Parameters @@ -602,7 +605,7 @@ class Geometry: return sorted(cells, key=lambda x: x.id) - def get_universes_by_name(self, name, case_sensitive=False, matching=False): + def get_universes_by_name(self, name, case_sensitive=False, matching=False) -> typing.List[openmc.Universe]: """Return a list of universes with matching names. Parameters @@ -623,7 +626,7 @@ class Geometry: """ return self._get_domains_by_name(name, case_sensitive, matching, 'universe') - def get_lattices_by_name(self, name, case_sensitive=False, matching=False): + def get_lattices_by_name(self, name, case_sensitive=False, matching=False) -> typing.List[openmc.Lattice]: """Return a list of lattices with matching names. Parameters @@ -644,7 +647,7 @@ class Geometry: """ return self._get_domains_by_name(name, case_sensitive, matching, 'lattice') - def remove_redundant_surfaces(self): + def remove_redundant_surfaces(self) -> dict: """Remove and return all of the redundant surfaces. Uses surface_precision attribute of Geometry instance for rounding and @@ -707,7 +710,7 @@ class Geometry: # Recursively traverse the CSG tree to count all cell instances self.root_universe._determine_paths(instances_only=instances_only) - def clone(self): + def clone(self) -> Geometry: """Create a copy of this geometry with new unique IDs for all of its enclosed materials, surfaces, cells, universes and lattices."""