Merge pull request #479 from opencybersecurityalliance/k2-session-gen…

…erator

add `execute_to_generate()` in Session

packages-nextgen/kestrel_core/src/kestrel/session.py

@@ -42,12 +42,24 @@ def execute(self, huntflow_block: str) -> Iterable[DataFrame]:
         huntflow blocks in this session, so all existing variables can be
         referred in the new huntflow block.
 
+        Parameters:
+            huntflow_block: the new huntflow block to be executed
+
+        Returns:
+            Evaluated result per Return instruction
+        """
+        return list(self.execute_to_generate(huntflow_block))
+
+    def execute_to_generate(self, huntflow_block: str) -> Iterable[DataFrame]:
+        """Execute a Kestrel huntflow and put results in a generator.
+
         Parameters:
             huntflow_block: the new huntflow block to be executed
 
         Yields:
-            Evaluated result per Return instruction in the huntflow block
+            Evaluated result per Return instruction
         """
+
         # TODO: return type generalization
 
         irgraph_new = parse_kestrel(huntflow_block)

packages-nextgen/kestrel_core/tests/test_session.py

@@ -21,7 +21,7 @@ def test_execute_in_cache():
                    ])
     b2 = DataFrame([ {"pid": 123} ])
     with Session() as session:
-        res = session.execute(hf)
+        res = session.execute_to_generate(hf)
         assert b1.equals(next(res))
         assert b2.equals(next(res))
         with pytest.raises(StopIteration):

packages-nextgen/kestrel_jupyter/src/kestrel_jupyter_kernel/kernel.py

@@ -34,7 +34,7 @@ def do_execute(
     ):
         if not silent:
             try:
-                for result in self.kestrel_session.execute(code):
+                for result in self.kestrel_session.execute_to_generate(code):
                     self.send_response(
                         self.iopub_socket,
                         "display_data",