adding plasma attributes

src/paramak/assemblies/spherical_tokamak.py

@@ -111,7 +111,7 @@ def spherical_tokamak_from_plasma(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {},
-):
+) -> Assembly:
     """Creates a spherical tokamak fusion reactor from a radial build and plasma parameters.
 
 
@@ -169,7 +169,7 @@ def spherical_tokamak(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {},
-):
+) -> Assembly:
     """  Creates a spherical tokamak fusion reactor from a radial build and vertical build.
 
     Args:
@@ -287,4 +287,6 @@ def spherical_tokamak(
 
     my_assembly.add(plasma, name="plasma", color=cq.Color(*colors.get("plasma", (0.5,0.5,0.5))))
 
+    my_assembly.elongation = elongation
+    my_assembly.triangularity = triangularity
     return my_assembly

src/paramak/assemblies/tokamak.py

@@ -34,7 +34,7 @@ def create_center_column_shield_cylinders(radial_build, rotation_angle, center_c
 
     number_of_cylinder_layers = count_cylinder_layers(radial_build)
 
-    for index, item in enumerate(radial_build):
+    for _, item in enumerate(radial_build):
         if item[0] == LayerType.PLASMA:
             break
 
@@ -158,7 +158,7 @@ def tokamak_from_plasma(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {}
-):
+) -> Assembly:
     """
     Creates a tokamak fusion reactor from a radial build and plasma parameters.
 
@@ -217,7 +217,7 @@ def tokamak(
     extra_cut_shapes: Sequence[cq.Workplane] = [],
     extra_intersect_shapes: Sequence[cq.Workplane] = [],
     colors: dict = {}
-):
+) -> Assembly:
     """
     Creates a tokamak fusion reactor from a radial and vertical build.
 

tests/test_assemblies/test_spherical_tokamak.py

@@ -133,4 +133,26 @@ def test_colors():
             "layer_4": (0.4, 0.4, 0.8),
             "layer_5": (0.5, 0.5, 0.8),
         },
-    )
+    )
+
+def test_attributes():
+    "passing in the colors dictionary should not raise an error"
+    my_reactor = paramak.spherical_tokamak_from_plasma(
+        radial_build=[
+            (paramak.LayerType.GAP, 10),
+            (paramak.LayerType.SOLID, 50),
+            (paramak.LayerType.SOLID, 15),
+            (paramak.LayerType.GAP, 50),
+            (paramak.LayerType.PLASMA, 300),
+            (paramak.LayerType.GAP, 60),
+            (paramak.LayerType.SOLID, 15),
+            (paramak.LayerType.SOLID, 60),
+            (paramak.LayerType.SOLID, 10),
+        ],
+        elongation=2,
+        triangularity=0.55,
+        rotation_angle=180,
+    )
+
+    assert my_reactor.elongation == 2
+    assert my_reactor.triangularity == 0.55