Fix zip operation to correctly interpret code gen results

Author: BenGalewsky

servicex_codegen/code_generator.py

@@ -32,10 +32,7 @@
 # modification, are permitted provided that the following conditions are met:
 #
 from abc import ABC, abstractmethod
-import os
-import zipfile
 from collections import namedtuple
-from tempfile import TemporaryDirectory
 
 GeneratedFileResult = namedtuple('GeneratedFileResult', 'hash output_dir')
 
@@ -48,41 +45,7 @@ def __init__(self, message: str):
 
 
 class CodeGenerator(ABC):
-    def zipdir(self, path: str, zip_handle: zipfile.ZipFile) -> None:
-        """Given a `path` to a directory, zip up its contents into a zip file.
-
-        Arguments:
-            path        Path to a local directory. The contents will be put into the zip file
-            zip_handle  The zip file handle to write into.
-        """
-        for root, _, files in os.walk(path):
-            for file in files:
-                zip_handle.write(os.path.join(root, file), file)
 
     @abstractmethod
-    def generate_code(self,  query, cache_path: str):
+    def generate_code(self, query, cache_path: str):
         pass
-
-    def translate_query_to_zip(self, query: str) -> bytes:
-        """Translate a text ast into a zip file as a memory stream
-
-        Arguments:
-            code            Text `qastle` version of the input ast generated by func_adl
-
-        Returns
-            bytes       Data that if written as a binary output would be a zip file.
-        """
-
-        # Generate the python code
-        with TemporaryDirectory() as tempdir:
-            r = self.generate_code(query, tempdir)
-
-            # Zip up everything in the directory - we are going to ship it as back as part
-            # of the message.
-            z_filename = os.path.join(str(tempdir), 'joined.zip')
-            zip_h = zipfile.ZipFile(z_filename, 'w', zipfile.ZIP_DEFLATED)
-            self.zipdir(r.output_dir, zip_h)
-            zip_h.close()
-
-            with open(z_filename, 'rb') as b_in:
-                return b_in.read()

servicex_codegen/post_operation.py

@@ -31,11 +31,13 @@
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
 #
+import os
+import zipfile
 from tempfile import TemporaryDirectory
 
 from flask import request, Response
 from flask_restful import Resource
-from servicex_codegen.code_generator import CodeGenerator
+from servicex_codegen.code_generator import CodeGenerator, GeneratedFileResult
 
 
 class GeneratedCode(Resource):
@@ -44,12 +46,47 @@ def make_api(cls, code_generator: CodeGenerator):
         cls.code_generator = code_generator
         return cls
 
+    def zipdir(self, path: str, zip_handle: zipfile.ZipFile) -> None:
+        """Given a `path` to a directory, zip up its contents into a zip file.
+
+        Arguments:
+            path        Path to a local directory. The contents will be put into the zip file
+            zip_handle  The zip file handle to write into.
+        """
+        for root, _, files in os.walk(path):
+            for file in files:
+                zip_handle.write(os.path.join(root, file), file)
+
+    def stream_generated_code(self, generated_code_result: GeneratedFileResult) -> bytes:
+        """Translate a text ast into a zip file as a memory stream
+
+        Arguments:
+            code            Text `qastle` version of the input ast generated by func_adl
+
+        Returns
+            bytes       Data that if written as a binary output would be a zip file.
+        """
+
+        # Generate the python code
+        with TemporaryDirectory() as tempdir:
+
+            # Zip up everything in the directory - we are going to ship it as back as part
+            # of the message.
+            z_filename = os.path.join(str(tempdir), 'joined.zip')
+            zip_h = zipfile.ZipFile(z_filename, 'w', zipfile.ZIP_DEFLATED)
+            self.zipdir(generated_code_result.output_dir, zip_h)
+            zip_h.close()
+
+            with open(z_filename, 'rb') as b_in:
+                return b_in.read()
+
     def post(self):
         try:
             with TemporaryDirectory() as tempdir:
                 code = request.data.decode('utf8')
-                zip_data = self.code_generator.generate_code(code, cache_path=tempdir)
+                generated_code_result = self.code_generator.generate_code(code, cache_path=tempdir)
 
+                zip_data = self.stream_generated_code(generated_code_result)
                 # Send the response back to you-know-what.
                 response = Response(
                     response=zip_data,

tests/test_code_generator.py

@@ -25,27 +25,53 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+import io
+import os
+import zipfile
+from tempfile import TemporaryDirectory
 
 from servicex_codegen import create_app
-from servicex_codegen.code_generator import GenerateCodeException
+from servicex_codegen.code_generator import GenerateCodeException, GeneratedFileResult
+
+
+def get_zipfile_data(zip_data: bytes):
+    with zipfile.ZipFile(io.BytesIO(zip_data)) as thezip:
+        for zipinfo in thezip.infolist():
+            with thezip.open(zipinfo) as thefile:
+                yield zipinfo.filename, thefile
+
+
+def check_zip_file(zip_data: bytes, expected_file_count):
+    names = []
+    for name, data in get_zipfile_data(zip_data):
+        names.append(name)
+        print(name)
+    assert len(names) == expected_file_count
 
 
 class TestPostOperation:
     def test_post_good_query(self, mocker):
         """Produce code for a simple good query"""
 
-        mock_ast_translator = mocker.Mock()
-        mock_ast_translator.generate_code = mocker.Mock(return_value="hi")
+        with TemporaryDirectory() as tempdir, \
+                open(os.path.join(tempdir, "baz.txt"), 'w'),\
+                open(os.path.join(tempdir, "foo.txt"), 'w'):
 
-        config = {
-            'TARGET_BACKEND': 'uproot'
-        }
-        app = create_app(config, provided_translator=mock_ast_translator)
-        client = app.test_client()
-        select_stmt = "(call ResultTTree (call Select (call SelectMany (call EventDataset (list 'localds://did_01')"  # noqa: E501
+            mock_ast_translator = mocker.Mock()
+            mock_ast_translator.generate_code = mocker.Mock(
+                return_value=GeneratedFileResult(hash="1234", output_dir=tempdir)
+            )
 
-        response = client.post("/servicex/generated-code", data=select_stmt)
+            config = {
+                'TARGET_BACKEND': 'uproot'
+            }
+            app = create_app(config, provided_translator=mock_ast_translator)
+            client = app.test_client()
+            select_stmt = "(call ResultTTree (call Select (call SelectMany (call EventDataset (list 'localds://did_01')"  # noqa: E501
+
+            response = client.post("/servicex/generated-code", data=select_stmt)
         assert response.status_code == 200
+        check_zip_file(response.data, 2)
         # Capture the temporary directory that was generated
         cache_dir = mock_ast_translator.generate_code.call_args[1]['cache_path']
         mock_ast_translator.generate_code.assert_called_with(select_stmt,