Added custom LofarDeviceProxy to enhance support for high-dimensional arrays in attribtues

lofar_station_client/devices.py

+""" Enhanced interfaces towards the station's Tango devices. """
+
+import ast
+from functools import lru_cache
+
+import numpy
+
+from tango import DeviceProxy
+from tango import ExtractAs
+
+
+class LofarDeviceProxy(DeviceProxy):
+    """A LOFAR-specific tango.DeviceProxy that provides
+    a richer experience."""
+
+    @lru_cache()
+    def get_attribute_config(self, name):
+        """Get cached attribute configurations, as they are not expected to change."""
+
+        return super().get_attribute_config(name)
+
+    @lru_cache()
+    def get_attribute_shape(self, name):
+        """Get the shape of the requested attribute, as a tuple."""
+
+        config = self.get_attribute_config(name)
+
+        if config.format and config.format[0] == "(":
+            # For >2D arrays, the "format" property describes actual
+            # the dimensions as a tuple (x, y, z, ...),
+            # so reshape the value accordingly.
+            shape = ast.literal_eval(config.format)
+        elif config.max_dim_y > 0:
+            # 2D array
+            shape = (config.max_dim_y, config.max_dim_x)
+        elif config.max_dim_x > 1:
+            # 1D array
+            shape = (config.max_dim_x,)
+        else:
+            # scalar
+            shape = ()
+
+        return shape
+
+    def read_attribute(self, name, extract_as=ExtractAs.Numpy):
+        """Read an attribute from the server."""
+
+        attr = super().read_attribute(name, extract_as)
+
+        # convert non-scalar values into their actual shape
+        shape = self.get_attribute_shape(name)
+        if shape != ():
+            attr.value = attr.value.reshape(shape)
+
+        return attr
+
+    def write_attribute(self, name, value):
+        """Write an attribute to the server."""
+
+        config = self.get_attribute_config(name)
+        shape = self.get_attribute_shape(name)
+
+        # 2D arrays also represent arrays of higher dimensionality. reshape them
+        # to fit their original Tango shape before writing.
+        if shape != ():
+            value = numpy.array(value)
+
+            if value.shape != shape:
+                raise ValueError(
+                    f"Invalid shape. Given: {value.shape} Expected: {shape}"
+                )
+
+            if len(shape) > 2:
+                # >2D arrays collapse into 2D
+                value = value.reshape((config.max_dim_y, config.max_dim_x))
+
+        return super().write_attribute(name, value)