📝 Add missing type hints

picometer/atom.py

@@ -51,15 +51,15 @@ def __init__(self,
         self.base = bf
         self.table = table
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.table) if self.table is not None else 0
 
-    def __add__(self, other):
+    def __add__(self, other) -> 'AtomSet':
         if not (self.base or self.table):
             return other
         elif not (other.base or other.table):
             return self
-        return AtomSet(self.base, pd.concat([self.table, other.table], axis=0))
+        return self.__class__(self.base, pd.concat([self.table, other.table], axis=0))
 
     def __getitem__(self, item) -> 'AtomSet':
         return self.__class__(bf=self.base, table=self.table[item])
@@ -91,11 +91,11 @@ def from_cif(cls, cif_path: str, block_name: str = None) -> 'AtomSet':
         return AtomSet(bf, atoms)
 
     @property
-    def fract_xyz(self):
+    def fract_xyz(self) -> np.ndarray:
         return np.vstack([self.table['fract_' + k].to_numpy() for k in 'xyz'])
 
     @property
-    def cart_xyz(self):
+    def cart_xyz(self) -> np.ndarray:
         return self.orthogonalise(self.fract_xyz)
 
     def fractionalise(self, cart_xyz: np.ndarray) -> np.ndarray:
@@ -140,7 +140,7 @@ def transform(self, symm_op_code: str) -> 'AtomSet':
         return self.__class__(self.base, data)
 
     @property
-    def centroid(self):
+    def centroid(self) -> np.ndarray:
         """A 3-vector with average atom position."""
         return self.cart_xyz.T.mean(axis=0)
 
@@ -169,7 +169,7 @@ def origin(self) -> Vector3:
         return self.centroid
 
     @origin.setter
-    def origin(self, new_origin):
+    def origin(self, new_origin) -> None:
         """Change origin to the new one provided in cartesian coordinates"""
         new_origin_fract = self.fractionalise(new_origin)
         delta = new_origin_fract - self.fractionalise(self.centroid)
@@ -202,7 +202,7 @@ def _distance(self, other: 'Shape') -> float:
             along = np.abs(np.dot(deltas, other.direction))
             return min(norms ** 2 - along ** 2)
 
-    def dihedral(self, *others: 'AtomSet'):
+    def dihedral(self, *others: 'AtomSet') -> float:
         assert all(o.kind is o.Kind.spatial for o in [self, *others])
         combined = sum(others, self)
         xyz = combined.cart_xyz.T

picometer/logging.py

@@ -30,11 +30,11 @@ def add_file_handler(path: Union[str, Path]) -> logging.FileHandler:
 
 class LogEventHandler(logging.Handler):
     """Custom handler for external log processing, see `register_log_listener`"""
-    def __init__(self, log_callback):
+    def __init__(self, log_callback) -> None:
         super().__init__()
         self.log_callback = log_callback  # A callback to process log events
 
-    def emit(self, record):
+    def emit(self, record) -> None:
         log_entry = self.format(record)
         self.log_callback(log_entry)  # Send log entry to the callback
 

picometer/models.py

@@ -19,7 +19,7 @@ def __init__(self,
         logger.debug(f'Initialized {self}')
 
     @property
-    def nodes(self):
+    def nodes(self) -> AtomSet:
         return self.atoms + self.centroids