Merge pull request #24 from sciencecorp/fix/stream-in-shape

sciencecorp · Oct 13, 2024 · 917d5fa · 917d5fa
2 parents d9d4676 + 2625dc0
commit 917d5fa
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/synapse-api b/synapse-api
diff --git a/synapse/client/nodes/stream_in.py b/synapse/client/nodes/stream_in.py
@@ -14,12 +14,13 @@
 class StreamIn(Node):
     type = NodeType.kStreamIn
 
-    def __init__(self, data_type: DataType):
+    def __init__(self, data_type: DataType, shape: List[int]):
         self.__sequence_number = 0
         self.__socket = socket.socket(
             socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP
         )
         self.data_type = data_type
+        self.shape = shape
 
     def write(self, data: SynapseData):
         if self.device is None:
@@ -76,6 +77,7 @@ def _to_proto(self):
         n = NodeConfig()
         i = StreamInConfig()
         i.data_type = self.data_type
+        i.shape = self.shape
 
         n.stream_in.CopyFrom(i)
         return n
@@ -88,4 +90,4 @@ def _from_proto(proto: Optional[StreamInConfig]):
         if not isinstance(proto, StreamInConfig):
             raise ValueError("proto is not of type StreamInConfig")
 
-        return StreamIn(data_type=proto.data_type)
+        return StreamIn(data_type=proto.data_type, shape=proto.shape)