Fix Pylint issues

ritchie46 · Oct 17, 2024 · 304199f · 304199f
1 parent 6d5715f
commit 304199f
Showing 1 changed file with 22 additions and 21 deletions.
diff --git a/anastruct/fem/system.py b/anastruct/fem/system.py
@@ -2291,29 +2291,31 @@ def apply_load_case(self, loadcase: LoadCase) -> None:
             kwargs = re.sub(r".??(\w+).?:", r"\1=", kwargs)
 
             exec(f"self.{method}({kwargs})")  # pylint: disable=exec-used
-    def get_stiffness_matrix(self, element_id):
-            """
-            Print the stiffness matrix for a specific element by its ID.
-            Args:
-                element_id (int): ID of the element.
-            """
-            try:
-                element = self.element_map[element_id]
-            except KeyError:
-                print(f"Element with ID {element_id} does not exist.")
-                return None
-
-            if isinstance(element, Element):
-                if hasattr(element, 'stiffness_matrix'):
-                    print(f"Stiffness Matrix for Element ID {element_id}:")
-                    return element.stiffness_matrix
-
-                print(f"Element ID {element_id} does not have a stiffness matrix.")
-            else:
-                print(f"Invalid element type for element ID {element_id}.")
+
+    def get_stiffness_matrix(self, element_id: int) -> Optional[Union[list, None]]:
+        """
+        Return the stiffness matrix for a specific element by its ID.
+        Args:
+            element_id (int): ID of the element.
+
+        Returns:
+        Optional[Union[list, None]]: The stiffness matrix of the element if it exists, otherwise None.
+        """
+        try:
+            element = self.element_map[element_id]
+        except KeyError:
+            print(f"Element with ID {element_id} does not exist.")
             return None
 
+        if isinstance(element, Element):
+            if hasattr(element, 'stiffness_matrix'):
+                print(f"Stiffness Matrix for Element ID {element_id}:")
+                return element.stiffness_matrix
 
+            print(f"Element ID {element_id} does not have a stiffness matrix.")
+        else:
+            print(f"Invalid element type for element ID {element_id}.")
+        return None
 
     def __deepcopy__(self, _: str) -> "SystemElements":
         """Deepcopy the SystemElements object.
@@ -2337,7 +2339,6 @@ def __deepcopy__(self, _: str) -> "SystemElements":
 
         return system
 
-
 def _negative_index_to_id(idx: int, collection: Collection[int]) -> int:
     """Convert a negative index to a positive index. (That is, allowing the Pythonic negative indexing)