L2SS-630: refactor statistics_client

tangostationcontrol/tangostationcontrol/clients/statistics_client.py

@@ -97,14 +97,7 @@ class StatisticsClient(AsyncCommClient):
         # redirect to right object. this works as long as the parameter names are unique among them.
         if annotation["type"] == "statistics":
             def read_function():
-                if annotation.get("reshape", False):
-                    # force array into the shape of the attribute
-                    if attribute.dim_y > 1:
-                        return self.collector.parameters[parameter].reshape(attribute.dim_y, attribute.dim_x)
-                    else:
-                        return self.collector.parameters[parameter].reshape(attribute.dim_x)
-                else:
-                    return self.collector.parameters[parameter]
+                process_statistics_annotation(self, annotation, attribute, parameter)
         elif annotation["type"] == "udp":
             def read_function():
                 return self.udp.parameters[parameter]
@@ -118,21 +111,26 @@ class StatisticsClient(AsyncCommClient):
             else:
                 raise ValueError(f"Unknown queue parameter requested: {parameter}")
         elif annotation["type"] == "replicator":
-            if parameter == "clients":
-                def read_function():
-                    return numpy.array(self.tcp.clients(),dtype=numpy.str)
-            elif parameter == "nof_bytes_sent":
-                def read_function():
-                    return numpy.uint64(self.tcp.nof_bytes_sent)
-            elif parameter == "nof_packets_sent":
-                def read_function():
-                    return numpy.uint64(self.tcp.nof_packets_sent)
-            elif parameter == "nof_tasks_pending":
+            parameters_dict = {"clients":           numpy.array(self.tcp.clients(),dtype=numpy.str),
+                            "nof_bytes_sent":       numpy.uint64(self.tcp.nof_bytes_sent),
+                            "nof_packets_sent":     numpy.uint64(self.tcp.nof_packets_sent),
+                            "nof_tasks_pending":    numpy.uint64(self.tcp.nof_tasks_pending)}
+            for k,v in parameters_dict.items():
+                if parameter == k:
                     def read_function():
-                    return numpy.uint64(self.tcp.nof_tasks_pending)
-            else:
+                        return v
             raise ValueError(f"Unknown replicator parameter requested: {parameter}")
 
+        def process_statistics_annotation(annotation, attribute, parameter):
+            if annotation.get("reshape", False):
+                # force array into the shape of the attribute
+                if attribute.dim_y > 1:
+                    return self.collector.parameters[parameter].reshape(attribute.dim_y, attribute.dim_x)
+                else:
+                    return self.collector.parameters[parameter].reshape(attribute.dim_x)
+            else:
+                return self.collector.parameters[parameter]
+        
         def write_function(value):
             """
             Not used here
@@ -140,3 +138,4 @@ class StatisticsClient(AsyncCommClient):
             pass
 
         return read_function, write_function   
+