From fb8c66a79aba63127ed2ae01c83365a4ad9eb6f0 Mon Sep 17 00:00:00 2001
From: shimwell <mail@jshimwell.com>
Date: Sun, 15 Dec 2024 02:09:42 +0100
Subject: [PATCH] adding plasma attributes

---
 src/paramak/assemblies/spherical_tokamak.py   |  6 +++--
 src/paramak/assemblies/tokamak.py             |  6 ++---
 .../test_assemblies/test_spherical_tokamak.py | 24 ++++++++++++++++++-
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/src/paramak/assemblies/spherical_tokamak.py b/src/paramak/assemblies/spherical_tokamak.py
index dad74678..6ef7d279 100644
--- a/src/paramak/assemblies/spherical_tokamak.py
+++ b/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
diff --git a/src/paramak/assemblies/tokamak.py b/src/paramak/assemblies/tokamak.py
index 29121929..68c38a9e 100644
--- a/src/paramak/assemblies/tokamak.py
+++ b/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.
 
diff --git a/tests/test_assemblies/test_spherical_tokamak.py b/tests/test_assemblies/test_spherical_tokamak.py
index ac4d0826..7ba12281 100644
--- a/tests/test_assemblies/test_spherical_tokamak.py
+++ b/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),
         },
-    )
\ No newline at end of file
+    )
+
+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
\ No newline at end of file