added support for numpy arrays

morioc/pvinfo.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 
 #PV data type: enum, string, char, float or int
 DTYPES = {
@@ -14,14 +16,32 @@ class PVInfo:
 
     def __init__(self, name, data, parse_string=True):
         self.name = name.upper()
-        self.dtype, self.value = infer_type(data, parse_string=parse_string)
+        #TODO: the following is ugly
+        res = infer_type(data, parse_string=parse_string)
+        if len(res) == 2:
+            self.dtype, self.value = res
+            self.count = 1
+        elif len(res) == 3:
+            self.dtype, self.value, self.count = res
 
-    def to_dict(self): #TODO count for arrays
-        return {"type": self.dtype}
+    def to_dict(self):
+        #TODO: the following is ugly
+        if self.count == 1:
+            return {"type": self.dtype}
+        else:
+            return {"type": self.dtype, "count": self.count}
 
 
 
-def infer_type(value, parse_string=True): #TODO arrays? strings?
+def infer_type(value, parse_string=True): #TODO lists/tuples/...? char vs. strings?
+    if isinstance(value, np.ndarray):
+        return infer_type_numpy(value) #TODO: also parse strings for np arrays?!
+    else:
+        return infer_type_scalar(value, parse_string=parse_string)
+
+
+
+def infer_type_scalar(value, parse_string=True):
     if parse_string and isinstance(value, str):
         if value in STR_REPS:
             value = STR_REPS[value]
@@ -32,12 +52,14 @@ def infer_type(value, parse_string=True): #TODO arrays? strings?
     else:
         dtype = type(value)
 
+    #TODO: this can probably be removed
     if isinstance(dtype, str):
         raise Exception("how?")
 
     if dtype in DTYPES:
         dtype = DTYPES[dtype]
     else:
+        # if dtype not understood, default to string and use string representation
         dtype = "string"
         value = str(value)
 
@@ -45,3 +67,20 @@ def infer_type(value, parse_string=True): #TODO arrays? strings?
 
 
 
+def infer_type_numpy(arr):
+    if arr.dtype == object:
+        raise NotImplementedError("cannot infer a single type from object arrays")
+
+    # bool -> int
+    if arr.dtype == bool:
+        arr = arr.astype(int)
+
+    # epics only allows 1D, and needs conversion to built-in data types
+    lst = arr.ravel().tolist()
+    dtype, _ = infer_type_scalar(lst[0])
+    count = len(lst)
+
+    return dtype, lst, count
+
+
+

morioc/test_infer_type.py

@@ -2,6 +2,8 @@
 
 from functools import partial
 from math import nan
+import numpy as np
+import pytest
 
 
 from pvinfo import infer_type as it
@@ -60,11 +62,48 @@ def test_psfalse_int():
     assert psfalse("int") == ("string", "int")
 
 
-def test_it_nan():
-    assert it(nan) == ("float", nan)
-
 def test_it_None():
     assert it(None) == ("string", "None")
 
+def test_it_True():
+    assert it(True) == ("string", "True")
+
+def test_it_False():
+    assert it(False) == ("string", "False")
+
+def test_it_nan():
+    assert it(nan) == ("float", nan)
+
+def test_it_np_nan():
+    assert it(np.nan) == ("float", np.nan)
+
+
+def test_it_np1D_int():
+    n = 10
+    arr = np.arange(n)
+    ref = arr.tolist()
+    assert it(arr) == ("int", ref, n)
+
+def test_it_np2D_int():
+    x = y = 10
+    n = x * y
+    arr = np.arange(n)
+    ref = arr.tolist()
+    arr = arr.reshape(x, y)
+    assert it(arr) == ("int", ref, n)
+
+def test_it_np1D_bool():
+    arr = np.array([True, False, True])
+    ref = [1, 0, 1]
+    assert it(arr) == ("int", ref, len(ref))
+
+def test_it_np1D_object():
+    arr = np.array([None]) # gives dtype=object
+    with pytest.raises(NotImplementedError):
+        it(arr)
+    arr = np.array(["test", 123, 1.23], dtype=object)
+    with pytest.raises(NotImplementedError):
+        it(arr)
+