fix mypy

materialsproject · Jul 31, 2023 · 185d2cb · 185d2cb
1 parent 28a895f
commit 185d2cb
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/pymatgen/core/structure.py b/pymatgen/core/structure.py
@@ -332,8 +332,9 @@ def __contains__(self, site: object) -> bool:
     def __iter__(self) -> Iterator[Site]:
         return self.sites.__iter__()
 
-    def __getitem__(self, ind):
-        return self.sites[ind]
+    # TODO return type needs fixing (can be list[Site] but raises lots of mypy errors)
+    def __getitem__(self, ind: int | slice) -> Site:
+        return self.sites[ind]  # type: ignore[return-value]
 
     def __len__(self) -> int:
         return len(self.sites)
@@ -1375,7 +1376,8 @@ def get_distance(self, i: int, j: int, jimage=None) -> float:
         Returns:
             distance
         """
-        return self[i].distance(self[j], jimage)
+        site: PeriodicSite = self[i]
+        return site.distance(self[j], jimage)
 
     def get_sites_in_sphere(
         self,