Lots of typehints

trussme/evaluate.py

@@ -1,16 +1,29 @@
+from typing import TypedDict
 import numpy
 
 
-def the_forces(truss_info):
-    tj = numpy.zeros([3, numpy.size(truss_info["connections"], axis=1)])
-    w = numpy.array([numpy.size(truss_info["reactions"], axis=0),
-                     numpy.size(truss_info["reactions"], axis=1)])
-    dof = numpy.zeros([3*w[1], 3*w[1]])
-    deflections = numpy.ones(w)
+TrussInfo = TypedDict("TrussInfo", {
+    "coordinates": numpy.ndarray,
+    "connections": numpy.ndarray,
+    "loads": numpy.ndarray,
+    "reactions": numpy.ndarray,
+    "area": float,
+    "elastic_modulus": float
+})
+
+
+def the_forces(truss_info: TrussInfo) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, float]:
+    tj: numpy.ndarray = numpy.zeros([3, numpy.size(truss_info["connections"], axis=1)])
+    w: numpy.ndarray = numpy.array([
+        numpy.size(truss_info["reactions"], axis=0),
+        numpy.size(truss_info["reactions"], axis=1)
+    ])
+    dof: numpy.ndarray = numpy.zeros([3*w[1], 3*w[1]])
+    deflections: numpy.ndarray = numpy.ones(w)
     deflections -= truss_info["reactions"]
 
     # This identifies joints that can be loaded
-    ff = numpy.where(deflections.T.flat == 1)[0]
+    ff: numpy.ndarray = numpy.where(deflections.T.flat == 1)[0]
 
     # Build the global stiffness matrix
     for i in range(numpy.size(truss_info["connections"], axis=1)):

trussme/joint.py

@@ -3,7 +3,7 @@
 
 class Joint(object):
 
-    def __init__(self, coordinates):
+    def __init__(self, coordinates: numpy.ndarray):
         # Save the joint id
         self.idx = -1
 
@@ -25,17 +25,17 @@ def __init__(self, coordinates):
         # Loads
         self.deflections = numpy.zeros([3, 1])
 
-    def free(self, d=3):
+    def free(self, d:int = 3):
         self.translation = numpy.zeros([3, 1])
         # If 2d, add out of plane support
         if d is 2:
             self.translation[2] = 1
 
-    def pinned(self, d=3):
+    def pinned(self, d: int = 3):
         # Restrict all translation
         self.translation = numpy.ones([3, 1])
 
-    def roller(self, axis='y', d=3):
+    def roller(self, axis: str = 'y', d: int = 3):
         # Only support reaction along denotated axis
         self.translation = numpy.zeros([3, 1])
         self.translation[ord(axis)-120] = 1

trussme/member.py

@@ -6,7 +6,7 @@
 class Member(object):
 
     # Shape types
-    shapes = ["pipe", "bar", "square", "box"]
+    shapes: list[str] = ["pipe", "bar", "square", "box"]
 
     def __init__(self, joint_a, joint_b):
         # Save id number