api client: handle stream decoding in swh.graph

swh/graph/client.py

@@ -22,34 +22,42 @@ class RemoteGraphClient(RPCClient):
         super().__init__(
             api_exception=GraphAPIError, url=url, timeout=timeout)
 
-    # Web API endpoints
+    def raw_verb_lines(self, verb, endpoint, **kwargs):
+        response = self.raw_verb(verb, endpoint, stream=True, **kwargs)
+        for line in response.iter_lines():
+            yield line.decode().lstrip('\n')
 
-    def leaves(self, src, edges="*", direction="forward"):
-        return self.get('leaves/{}'.format(src),
-                        params={
-                            'edges': edges,
-                            'direction': direction
-                        },
-                        stream_lines=True)
+    def get_lines(self, endpoint, **kwargs):
+        yield from self.raw_verb_lines('get', endpoint, **kwargs)
 
-    def neighbors(self, src, edges="*", direction="forward"):
-        return self.get('neighbors/{}'.format(src),
-                        params={
-                            'edges': edges,
-                            'direction': direction
-                        },
-                        stream_lines=True)
+    # Web API endpoints
 
     def stats(self):
         return self.get('stats')
 
+    def leaves(self, src, edges="*", direction="forward"):
+        return self.get_lines(
+            'leaves/{}'.format(src),
+            params={
+                'edges': edges,
+                'direction': direction
+            })
+
+    def neighbors(self, src, edges="*", direction="forward"):
+        return self.get_lines(
+            'neighbors/{}'.format(src),
+            params={
+                'edges': edges,
+                'direction': direction
+            })
+
     def visit_nodes(self, src, edges="*", direction="forward"):
-        return self.get('visit/nodes/{}'.format(src),
-                        params={
-                            'edges': edges,
-                            'direction': direction
-                        },
-                        stream_lines=True)
+        return self.get_lines(
+            'visit/nodes/{}'.format(src),
+            params={
+                'edges': edges,
+                'direction': direction
+            })
 
     def visit_paths(self, src, edges="*", direction="forward"):
         def decode_path_wrapper(it):
@@ -57,17 +65,18 @@ class RemoteGraphClient(RPCClient):
                 yield json.loads(e)
 
         return decode_path_wrapper(
-            self.get('visit/paths/{}'.format(src),
-                     params={
-                         'edges': edges,
-                         'direction': direction
-                     }, stream_lines=True))
+            self.get_lines(
+                'visit/paths/{}'.format(src),
+                params={
+                    'edges': edges,
+                    'direction': direction
+                }))
 
     def walk(self, src, dst, edges="*", traversal="dfs", direction="forward"):
-        return self.get('walk/{}/{}'.format(src, dst),
-                        params={
-                            'edges': edges,
-                            'traversal': traversal,
-                            'direction': direction
-                        },
-                        stream_lines=True)
+        return self.get_lines(
+            'walk/{}/{}'.format(src, dst),
+            params={
+                'edges': edges,
+                'traversal': traversal,
+                'direction': direction
+            })