From 4cc14b70f26092d699e4a028de58611f1f4b5346 Mon Sep 17 00:00:00 2001
From: BitterB0NG0 <93603537+BitterB0NG0@users.noreply.github.com>
Date: Sun, 27 Oct 2024 22:22:02 -0700
Subject: [PATCH] files/inp: added ``__getitem__`` methods

---
 src/pymcnp/files/inp/cells.py    | 9 +++++++++
 src/pymcnp/files/inp/surfaces.py | 9 +++++++++
 2 files changed, 18 insertions(+)

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