add representation methods and encoders

The encoders are useful to store objects on disk between multiple calls to a program. This is useful in some cases where you need to recover the previous state of your run. The representations implementation is very useful to print the results of the requests to know what you got. It is very useful for debugging and testing out things when developping. Signed-off-by: Arnaud Fiorini <[email protected]>
eclipse-cdt-cloud · Jul 10, 2024 · cfa5eea · cfa5eea
1 parent bae43f9
commit cfa5eea
Show file tree

Hide file tree

Showing 12 changed files with 288 additions and 64 deletions.
diff --git a/tsp/entry.py b/tsp/entry.py
@@ -21,6 +21,7 @@
 # SOFTWARE.
 
 """Entry classes file."""
+import json
 
 from tsp.output_element_style import OutputElementStyle
 
@@ -82,3 +83,27 @@ def __init__(self, params):
             if params.get(STYLE_KEY) is not None:
                 self.style = OutputElementStyle(params.get(STYLE_KEY))
             del params[STYLE_KEY]
+
+class EntryHeaderEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, EntryHeader):
+            return {'name': obj.name}
+        return super().default(obj)
+
+class EntryEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, Entry):
+            return {
+                'id': obj.id,
+                'parent_id': obj.parent_id,
+                'labels': obj.labels,
+                'style': EntryElementStyleEncoder().default(obj.style) if obj.style else None
+            }
+        return super().default(obj)
+
+class EntryElementStyleEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, OutputElementStyle):
+            # Assuming OutputElementStyle has a to_dict method
+            return obj.to_dict()
+        return super().default(obj)
diff --git a/tsp/entry_model.py b/tsp/entry_model.py
@@ -21,10 +21,11 @@
 # SOFTWARE.
 
 """EntryModel class file."""
+import json
 
 from tsp.model_type import ModelType
-from tsp.time_graph_model import TimeGraphEntry
-from tsp.entry import EntryHeader, Entry
+from tsp.time_graph_model import TimeGraphEntry, TimeGraphEntryEncoder
+from tsp.entry import EntryHeader, Entry, EntryHeaderEncoder, EntryEncoder
 
 HEADER_KEY = "headers"
 ENTRIES_KEY = "entries"
@@ -56,3 +57,15 @@ def __init__(self, params, model_type=ModelType.XY_TREE):
                 else:
                     self.entries.append(Entry(entry))
             del params[ENTRIES_KEY]
+
+    def __repr__(self) -> str:
+        return 'EntryModel({})'.format(', '.join(str(entry) for entry in self.entries))
+
+class EntryModelEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, EntryModel):
+            return {
+                'headers': [EntryHeaderEncoder().default(header) for header in obj.headers],
+                'entries': [TimeGraphEntryEncoder().default(entry) if isinstance(entry, TimeGraphEntry) else EntryEncoder().default(entry) for entry in obj.entries]
+            }
+        return super().default(obj)
diff --git a/tsp/experiment.py b/tsp/experiment.py
@@ -95,3 +95,8 @@ def __init__(self, params):
         # Array of all the traces contained in the experiment
         if TRACES_TIME_KEY in params:
             self.traces = TraceSet(params.get(TRACES_TIME_KEY))
+
+    def __repr__(self):
+        return 'Experiment({}: UUID={}, start={}, end={}, nevent={}, traces={}, indexing={})'.format(
+            self.name, self.UUID, self.start, self.end, self.number_of_events, self.traces, self.indexing_status
+        )
diff --git a/tsp/output_descriptor.py b/tsp/output_descriptor.py
@@ -22,6 +22,8 @@
 
 """OutputDescriptor class file."""
 
+import json
+
 NA = "N/A"
 UNKOWN = "UNKNOWN"
 ID_KEY = "id"
@@ -113,3 +115,22 @@ def __init__(self, params):
             del params[COMPATIBLE_PROVIDERS_KEY]
         else:
             self.compatible_providers = []
+
+    def __repr__(self):
+        return 'OutputDescriptor(id={}, name={}, description={})'.format(self.id, self.name, self.description)
+
+class OutputDescriptorEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, OutputDescriptor):
+            return {
+                'id': obj.id,
+                'name': obj.name,
+                'description': obj.description,
+                'type': obj.type,
+                'query_parameters': obj.query_parameters,
+                'start': obj.start,
+                'end': obj.end,
+                'final': obj.final,
+                'compatible_providers': obj.compatible_providers
+            }
+        return super().default(obj)
diff --git a/tsp/output_descriptor_set.py b/tsp/output_descriptor_set.py
@@ -38,3 +38,6 @@ def __init__(self, params):
         self.descriptors = []
         for obj in params:
             self.descriptors.append(OutputDescriptor(obj))
+
+    def __repr__(self):
+        return ', '.join([str(descriptor) for descriptor in self.descriptors])
diff --git a/tsp/response.py b/tsp/response.py
@@ -22,12 +22,15 @@
 
 """Response classes file."""
 
+import json
+
 from enum import Enum
 
 from tsp.model_type import ModelType
 from tsp.output_descriptor import OutputDescriptor
-from tsp.entry_model import EntryModel
-from tsp.xy_model import XYModel
+from tsp.entry_model import EntryModel, EntryModelEncoder
+from tsp.xy_model import XYModel, XYModelEncoder
+from tsp.time_graph_model import TimeGraphModel, TimeGraphArrow, TimeGraphModelEncoder, TimeGraphArrowEncoder
 
 MODEL_KEY = "model"
 OUTPUT_DESCRIPTOR_KEY = "output"
@@ -102,6 +105,25 @@ def __init__(self, params, model_type):
 
         # Message associated with the response
         if STATUS_MESSAGE_KEY in params:
-            self.status = params.get(STATUS_MESSAGE_KEY)
+            self.status_text = params.get(STATUS_MESSAGE_KEY)
         else:  # pragma: no cover
-            self.status = ""
+            self.status_text = ""
+
+
+class GenericResponseEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, GenericResponse):
+            model = obj.model
+            if obj.model_type ==  ModelType.TIME_GRAPH_TREE \
+                or obj.model_type == ModelType.XY_TREE \
+                or obj.model_type == ModelType.DATA_TREE:
+                model = EntryModelEncoder().default(obj.model)
+            elif obj.model_type == ModelType.TIME_GRAPH_STATE:
+                model = TimeGraphModelEncoder().default(obj.model)
+            elif obj.model_type == ModelType.TIME_GRAPH_ARROW:
+                model = [TimeGraphArrowEncoder().default(arrow) for arrow in obj.model]
+            elif obj.model_type == ModelType.XY:
+                model = XYModelEncoder().default(obj.model)
+            return model
+
+        return super().default(obj)
diff --git a/tsp/time_graph_model.py b/tsp/time_graph_model.py
@@ -21,8 +21,8 @@
 # SOFTWARE.
 
 """TimeGraph classes file."""
-
-from tsp.entry import Entry
+import json
+from tsp.entry import Entry, EntryElementStyleEncoder
 
 TYPE_KEY = "type"
 START_TIME_KEY = "startTime"
@@ -70,6 +70,8 @@ def __init__(self, params):
             self.has_row_model = params.get(HAS_ROW_MODEL_KEY)
             del params[HAS_ROW_MODEL_KEY]
 
+    def __repr__(self) -> str:
+        return 'TimeGraphEntry(start={}, end={})'.format(self.start_time, self.end_time)
 
 class TimeGraphModel:
     '''
@@ -83,6 +85,8 @@ def __init__(self, params):
                 self.rows.append(TimeGraphRow(row))
             del params[ROWS_KEY]
 
+    def __repr__(self) -> str:
+        return 'TimeGraphModel({})'.format(', '.join(str(row) for row in self.rows))
 
 class TimeGraphRow:
     '''
@@ -104,6 +108,9 @@ def __init__(self, params):
                 self.states.append(TimeGraphState(state))
             del params[STATES_KEY]
 
+    def __repr__(self) -> str:
+        return 'TimeGraphRow({})'.format(', '.join([str(state) for state in self.states]))
+
 
 class TimeGraphState:
     '''
@@ -143,6 +150,12 @@ def __init__(self, params):
             self.style = params.get(STYLE_KEY)
             del params[STYLE_KEY]
 
+    def __repr__(self) -> str:
+        representation = 'TimeGraphState('
+        if hasattr(self, 'label'):
+            representation += '{}, '.format(self.label)
+        return representation + 'start={}, end={})'.format(self.start_time, self.end_time)
+
 
 class TimeGraphArrow:
     '''
@@ -181,3 +194,59 @@ def __init__(self, params):
         if STYLE_KEY in params:
             self.style = params.get(STYLE_KEY)
             del params[STYLE_KEY]
+
+class TimeGraphEntryEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, TimeGraphEntry):
+            return {
+                'id': obj.id,
+                'parent_id': obj.parent_id,
+                'labels': obj.labels,
+                'style': EntryElementStyleEncoder().default(obj.style) if obj.style else None,
+                'start': obj.start_time,
+                'end': obj.end_time,
+                'hasData': obj.has_row_model
+            }
+        return super().default(obj)
+
+class TimeGraphModelEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, TimeGraphModel):
+            return {
+                'rows': [TimeGraphRowEncoder().default(row) for row in obj.rows]
+            }
+        return super().default(obj)
+
+class TimeGraphRowEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, TimeGraphRow):
+            return {
+                'entry_id': obj.entry_id,
+                'states': [TimeGraphStateEncoder().default(state) for state in obj.states]
+            }
+        return super().default(obj)
+
+class TimeGraphStateEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, TimeGraphState):
+            return {
+                'start': obj.start_time,
+                'end': obj.end_time,
+                # 'duration': obj.duration,
+                # 'values': obj.value,
+                # 'tags': obj.tags,
+                'style': obj.style
+            }
+        return super().default(obj)
+
+class TimeGraphArrowEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, TimeGraphArrow):
+            return {
+                'sourceId': obj.source_id,
+                'targetId': obj.target_id,
+                'start': obj.start,
+                'end': obj.end,
+                'style': obj.style
+            }
+        return super().default(obj)
diff --git a/tsp/trace.py b/tsp/trace.py
@@ -98,3 +98,8 @@ def __init__(self, params):
             del params[INDEXING_STATUS_KEY]
         else:  # pragma: no cover
             self.indexing_status = 0
+
+    def __repr__(self):
+        return 'Trace({}: UUID={}, start={}, end={}, nevent={}, path={}, indexing={})'.format(
+            self.name, self.UUID, self.start, self.end, self.number_of_events, self.path, self.indexing_status
+        )
diff --git a/tsp/trace_set.py b/tsp/trace_set.py
@@ -38,3 +38,6 @@ def __init__(self, params):
         self.traces = []
         for obj in params:
             self.traces.append(Trace(obj))
+
+    def __repr__(self) -> str:
+        return 'TraceSet({})'.format(', '.join(str(trace) for trace in self.traces))
diff --git a/tsp/tsp_client.py b/tsp/tsp_client.py
@@ -101,7 +101,7 @@ def fetch_trace(self, uuid):
             print("get trace failed: {0}".format(response.status_code))
             return TspClientResponse(None, response.status_code, response.text)
 
-    def open_trace(self, name, path):
+    def open_trace(self, path, name) -> TspClientResponse:
         '''
         Open a trace on the server
         parameters: Query object
@@ -113,11 +113,33 @@ def open_trace(self, name, path):
         my_parameters = {'name': name, 'uri': path}
         parameters = {'parameters': my_parameters}
 
+        response = requests.post(api_url, json=parameters, headers=headers)
+        response.raise_for_status()
+        return TspClientResponse(Trace(json.loads(response.content.decode('utf-8'))),
+                                 response.status_code, response.text)
+
+    def open_traces(self, path, name, max_depth, filter):
+        '''
+        Open a trace on the server
+        parameters: Query object
+        :return: :class:`TspClientResponse <Trace>` object
+        :rtype: TspClientResponse
+        '''
+        api_url = '{0}traces'.format(self.base_url)
+
+        my_parameters = {'name': name, 'uri': path, 'maxDepth': max_depth}
+        parameters = {'parameters': my_parameters}
+
         response = requests.post(api_url, json=parameters, headers=headers)
 
         if response.status_code == 200:
-            return TspClientResponse(Trace(json.loads(response.content.decode('utf-8'))),
-                                     response.status_code, response.text)
+            response_size = len(bytes(response.text, 'utf-8'))
+            content = json.loads(response.content.decode('utf-8'))
+            traces = []
+            for trace in content:
+                traces.append(Trace(trace))
+            return TspClientResponse(traces,
+                                     response.status_code, response.text, response_size)
         else:  # pragma: no cover
             print("post trace failed: {0}".format(response.status_code))
             return TspClientResponse(None, response.status_code, response.text)
@@ -281,7 +303,7 @@ def fetch_datatree(self, exp_uuid, output_id, parameters=None):
             print(GET_TREE_FAILED.format(response.status_code))
             return TspClientResponse(None, response.status_code, response.text)
 
-    def fetch_timegraph_tree(self, exp_uuid, output_id, parameters=None):
+    def fetch_timegraph_tree(self, exp_uuid, output_id, parameters=None) -> TspClientResponse:
         '''
         Fetch Time Graph tree, Model extends TimeGraphEntry
         :param exp_uuid: Experiment UUID
@@ -300,9 +322,10 @@ def fetch_timegraph_tree(self, exp_uuid, output_id, parameters=None):
         response = requests.post(api_url, json=params, headers=headers)
 
         if response.status_code == 200:
+            response_size = len(bytes(response.text, 'utf-8'))
             return TspClientResponse(GenericResponse(json.loads(response.content.decode('utf-8')),
                                                      ModelType.TIME_GRAPH_TREE),
-                                     response.status_code, response.text)
+                                     response.status_code, response.text, response_size)
         else:  # pragma: no cover
             print(GET_TREE_FAILED.format(response.status_code))
             return TspClientResponse(None, response.status_code, response.text)