Handle Vertex and VertexLike types consistently

anastruct/fem/system.py

@@ -19,15 +19,13 @@
 import numpy as np
 
 from anastruct.basic import FEMException, arg_to_list
-from anastruct.fem import plotter
+from anastruct.fem import plotter, system_components
 from anastruct.fem.elements import Element
 from anastruct.fem.postprocess import SystemLevel as post_sl
 from anastruct.fem.util.load import LoadCase
 from anastruct.sectionbase import properties
 from anastruct.vertex import Vertex, vertex_range
 
-from . import system_components
-
 if TYPE_CHECKING:
     from matplotlib.figure import Figure
 
@@ -39,6 +37,7 @@
         MpType,
         Spring,
         SupportDirection,
+        VertexLike,
     )
 
 
@@ -195,7 +194,7 @@ def id_last_node(self) -> int:
 
     def add_sequential_elements(
         self,
-        location: Union[Sequence[Sequence[float]], Sequence[Vertex]],
+        location: Sequence["VertexLike"],
         EA: Optional[Union[List[float], np.ndarray, float]] = None,
         EI: Optional[Union[List[float], np.ndarray, float]] = None,
         g: Optional[Union[List[float], np.ndarray, float]] = None,
@@ -206,7 +205,7 @@ def add_sequential_elements(
         """Add multiple elements based upon any number of sequential points.
 
         Args:
-            location (Union[Sequence[Sequence[float]], Sequence[Vertex]]): Sequence of points
+            location (Sequence[VertexLike]): Sequence of points
                 that define the elements.
             EA (Optional[Union[List[float], np.ndarray, float]], optional): Axial stiffnesses. Defaults to None.
             EI (Optional[Union[List[float], np.ndarray, float]], optional): Bending stiffnesses. Defaults to None.
@@ -300,9 +299,7 @@ def add_element_grid(
 
     def add_truss_element(
         self,
-        location: Union[
-            Sequence[Sequence[float]], Sequence[Vertex], Sequence[float], Vertex
-        ],
+        location: Union[Sequence["VertexLike"], "VertexLike"],
         EA: Optional[float] = None,
         **kwargs: dict,
     ) -> int:
@@ -316,7 +313,7 @@ def add_truss_element(
                    location=Vertex
 
         Args:
-            location (Union[ Sequence[Sequence[float]], Sequence[Vertex], Sequence[float], Vertex ]):
+            location (Union[Sequence[VertexLike], VertexLike]):
                 The two nodes of the element or the next node of the element.
             EA (Optional[float], optional): Axial stiffness of the new element. Defaults to None.
 
@@ -336,9 +333,7 @@ def add_truss_element(
 
     def add_element(
         self,
-        location: Union[
-            Sequence[Sequence[float]], Sequence[Vertex], Sequence[float], Vertex
-        ],
+        location: Union[Sequence["VertexLike"], "VertexLike"],
         EA: Optional[float] = None,
         EI: Optional[float] = None,
         g: float = 0,
@@ -363,7 +358,7 @@ def add_element(
                 spring={1: 0}
 
         Args:
-            location (Union[ Sequence[Sequence[float]], Sequence[Vertex], Sequence[float], Vertex ]):
+            location (Union[Sequence[VertexLike], VertexLike]):
                 The two nodes of the element or the next node of the element
             EA (Optional[float], optional): Axial stiffness of the new element. Defaults to None.
             EI (Optional[float], optional): Bending stiffness of the new element. Defaults to None.
@@ -477,9 +472,7 @@ def add_element(
 
     def add_multiple_elements(
         self,
-        location: Union[
-            Sequence[Sequence[float]], Sequence[Vertex], Sequence[float], Vertex
-        ],
+        location: Union[Sequence["VertexLike"], "VertexLike"],
         n: Optional[int] = None,
         dl: Optional[float] = None,
         EA: Optional[float] = None,
@@ -498,7 +491,7 @@ def add_multiple_elements(
                 last={'EA': 1e3, 'mp': 290}
 
         Args:
-            location (Union[ Sequence[Sequence[float]], Sequence[Vertex], Sequence[float], Vertex ]):
+            location (Union[Sequence[VertexLike], VertexLike]):
                 The two nodes of the element or the next node of the element.
             n (Optional[int], optional): Number of elements to add between the first and last nodes. Defaults to None.
             dl (Optional[float], optional): Length of sub-elements to add between the first and last nodes.
@@ -618,7 +611,7 @@ def add_multiple_elements(
     def insert_node(
         self,
         element_id: int,
-        location: Optional[Union[Sequence[float], Vertex]] = None,
+        location: Optional["VertexLike"] = None,
         factor: Optional[float] = None,
     ) -> None:
         """Insert a node into an existing structure.
@@ -631,7 +624,7 @@ def insert_node(
 
         Args:
             element_id (int): Id number of the element in which you want to insert the node
-            location (Optional[Union[Sequence[float], Vertex]], optional): Location in which to insert the node.
+            location (Optional[VertexLike], optional): Location in which to insert the node.
                 Defaults to None.
             factor (Optional[float], optional): Fraction of distance from start to end of elmeent on which to
                 divide the element. Must be between 0 and 1. Defaults to None.
@@ -1651,12 +1644,7 @@ def find_node_id(self, vertex: Union[Vertex, Sequence[float]]) -> Optional[int]:
         Returns:
             Optional[int]: id of the node at the location of the vertex
         """
-        if isinstance(vertex, (list, tuple)):
-            vertex_v = Vertex(vertex)
-        elif isinstance(vertex, Vertex):
-            vertex_v = vertex
-        else:
-            raise TypeError("vertex must be a list, tuple or Vertex")
+        vertex_v = Vertex(vertex)
         try:
             tol = 1e-9
             return next(

anastruct/fem/system_components/util.py

@@ -1,11 +1,14 @@
 from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union
 
+import numpy as np
+
 from anastruct.basic import FEMException, angle_x_axis
 from anastruct.fem.node import Node
 from anastruct.vertex import Vertex
 
 if TYPE_CHECKING:
     from anastruct.fem.system import MpType, Spring, SystemElements
+    from anastruct.types import VertexLike
 
 
 def check_internal_hinges(system: "SystemElements", node_id: int) -> None:
@@ -78,11 +81,8 @@ def append_node_id(
 def det_vertices(
     system: "SystemElements",
     location_list: Union[
-        Vertex,
-        Sequence[Vertex],
-        Sequence[int],
-        Sequence[float],
-        Sequence[Sequence[float]],
+        "VertexLike",
+        Sequence["VertexLike"],
     ],
 ) -> Tuple[Vertex, Vertex]:
     """Determine the vertices of a location list
@@ -101,23 +101,21 @@ def det_vertices(
     """
     if isinstance(location_list, Vertex):
         point_1 = system._previous_point
-        point_2 = Vertex(location_list)
-    elif len(location_list) == 1 and isinstance(location_list[0], Sequence):
+        point_2 = location_list
+    elif isinstance(location_list, Sequence) and len(location_list) == 1:
         point_1 = system._previous_point
-        point_2 = Vertex(location_list[0][0], location_list[0][1])
-    elif isinstance(location_list[0], (int, float)) and isinstance(
-        location_list[1], (int, float)
+        point_2 = Vertex(location_list[0])
+    elif (
+        isinstance(location_list, Sequence)
+        and len(location_list) == 2
+        and isinstance(location_list[0], (float, int, np.number))
+        and isinstance(location_list[1], (float, int, np.number))
     ):
         point_1 = system._previous_point
         point_2 = Vertex(location_list[0], location_list[1])
-    elif isinstance(location_list[0], Vertex) and isinstance(location_list[1], Vertex):
-        point_1 = location_list[0]
-        point_2 = location_list[1]
-    elif isinstance(location_list[0], Sequence) and isinstance(
-        location_list[1], Sequence
-    ):
-        point_1 = Vertex(location_list[0][0], location_list[0][1])
-        point_2 = Vertex(location_list[1][0], location_list[1][1])
+    elif isinstance(location_list, Sequence) and len(location_list) == 2:
+        point_1 = Vertex(location_list[0])
+        point_2 = Vertex(location_list[1])
     else:
         raise FEMException(
             "Flawed inputs",

anastruct/types.py

@@ -1,10 +1,17 @@
-from typing import Dict, Literal
+from typing import TYPE_CHECKING, Dict, Literal, Sequence, Union
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from anastruct.vertex import Vertex
 
 AxisNumber = Literal[1, 2, 3]
 Dimension = Literal["x", "y", "z", "both"]
 ElementType = Literal["general", "truss"]
 LoadDirection = Literal["element", "x", "y", "parallel", "perpendicular", "angle"]
 MpType = Dict[Literal[1, 2], float]
+NumberLike = Union[float, int, np.number]
 OrientAxis = Literal["y", "z"]
 Spring = Dict[Literal[1, 2], float]
 SupportDirection = Literal["x", "y", "1", "2", 1, 2]
+VertexLike = Union[Sequence[Union[float, int]], np.ndarray, "Vertex"]