From 03f75198c35ca82c5bd3f09824686f0c74c4a07d Mon Sep 17 00:00:00 2001
From: Christopher Mayes <31023527+ChristopherMayes@users.noreply.github.com>
Date: Tue, 19 Nov 2024 14:57:00 -0800
Subject: [PATCH] Fix nice array for single element list-like

---
 pmd_beamphysics/units.py |  1 +
 tests/test_nice.py       | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)
 create mode 100644 tests/test_nice.py

diff --git a/pmd_beamphysics/units.py b/pmd_beamphysics/units.py
index be8befa..1b06be6 100644
--- 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
new file mode 100644
index 0000000..25f04d2
--- /dev/null
+++ 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]))