diff --git a/src/pymcnp/files/inp/cells.py b/src/pymcnp/files/inp/cells.py index 4e6cac1..a4ffe00 100644 --- a/src/pymcnp/files/inp/cells.py +++ b/src/pymcnp/files/inp/cells.py @@ -80,3 +80,12 @@ def to_arguments(self) -> list: """ return [card.to_arguments() for card in self._cards.values()] + + def __getitem__(self, index: int) -> Cell: + try: + return self._cards[index] + except KeyError: + raise KeyError + + def __contains__(self, item: str) -> Cell: + return item in self._cards diff --git a/src/pymcnp/files/inp/surfaces.py b/src/pymcnp/files/inp/surfaces.py index 30ff8ff..9bba500 100644 --- a/src/pymcnp/files/inp/surfaces.py +++ b/src/pymcnp/files/inp/surfaces.py @@ -122,3 +122,12 @@ def to_cadquery_file(self, filename: str, hasHeader: bool = True) -> None: with open(filename, 'w') as file: file.write(self.to_cadquery(hasHeader)) + + def __getitem__(self, index: int) -> Surface: + try: + return self._cards[index] + except KeyError: + raise KeyError + + def __contains__(self, item: str) -> Surface: + return item in self._cards