Fix nice array for single element list-like

ChristopherMayes · Nov 19, 2024 · 03f7519 · 03f7519
1 parent 8be7b4e
commit 03f7519
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/pmd_beamphysics/units.py b/pmd_beamphysics/units.py
@@ -379,6 +379,7 @@ def nice_array(a):
     if np.isscalar(a):
         x = a
     elif len(a) == 1:
+        a = np.asarray(a)
         x = a[0]
     else:
         a = np.array(a)

diff --git a/tests/test_nice.py b/tests/test_nice.py
@@ -0,0 +1,17 @@
+from pmd_beamphysics.units import nice_array
+import numpy as np
+
+
+def test_nice_array():
+    nice_array(1)
+
+    nice_array([1])
+
+    nice_array(
+        [
+            1,
+            2,
+        ]
+    )
+
+    nice_array(np.array([1, 2]))