Add in-memory archiving option for production (#72)

* Add in-memory archiving for production

* Apply suggestions from review

Author: trsvchn

app/codegen.py

@@ -1,6 +1,10 @@
 """Code Generator base module.
 """
+import base64
+import io
 import shutil
+import tarfile
+import zipfile
 from pathlib import Path
 
 from jinja2 import Environment, FileSystemLoader
@@ -46,7 +50,7 @@ def make_and_write(self, template_name: str, dest_path: Path):
         for fname, code in self.rendered_code.items():
             (dest_path / template_name / fname).write_text(code)
 
-    def make_archive(self, template_name, archive_format, dest_path):
+    def write_archive(self, template_name, archive_format, dest_path):
         """Creates dist dir with generated code, then makes the archive."""
         self.make_and_write(template_name, dest_path)
         archive_fname = shutil.make_archive(
@@ -57,3 +61,26 @@ def make_archive(self, template_name, archive_format, dest_path):
         )
         archive_fname = shutil.move(archive_fname, dest_path / archive_fname.split("/")[-1])
         return archive_fname
+
+    def writes_archive(self, template_name, archive_format):
+        """Writes archive as Base64 string."""
+        arch_buffer = io.BytesIO()
+
+        if archive_format == "zip":
+            with zipfile.ZipFile(arch_buffer, "w", zipfile.ZIP_DEFLATED) as arch:
+                for fname, code in self.rendered_code.items():
+                    arch.writestr(f"{template_name}/{fname}", code)
+
+        elif archive_format == "tar.gz":
+            with tarfile.open(fileobj=arch_buffer, mode="w:gz") as arch:
+                for fname, code in self.rendered_code.items():
+                    tarinfo = tarfile.TarInfo(name=f"{template_name}/{fname}")
+                    code_fileobj = io.BytesIO(code.encode())
+                    tarinfo.size = len(code_fileobj.getvalue())
+                    arch.addfile(tarinfo, code_fileobj)
+        else:
+            raise ValueError(f"Wrong archive format '{archive_format}', use one of available formats: zip, tar.gz")
+
+        arch_str = base64.b64encode(arch_buffer.getvalue()).decode()
+
+        return arch_str

app/streamlit_app.py

@@ -1,3 +1,4 @@
+import os
 import shutil
 import tempfile
 from pathlib import Path
@@ -7,7 +8,7 @@
 from utils import import_from_file
 
 __version__ = "0.1.0"
-
+DEV_MODE = os.getenv("DEV_MODE", 0) == 1
 
 FOLDER_TO_TEMPLATE_NAME = {
     "Image Classification": "image_classification",
@@ -127,7 +128,7 @@ def add_content(self):
                 if len(code):  # don't show files which don't have content in them
                     self.render_code(fname, code)
 
-    def add_download(self):
+    def download_panel_dev(self):
         st.markdown("")
         col1, col2 = st.beta_columns(2)
         with col1:
@@ -136,27 +137,58 @@ def add_download(self):
             # https://github.com/streamlit/streamlit/issues/400
             # https://github.com/streamlit/streamlit/issues/400#issuecomment-648580840
             if st.button("Generate an archive"):
-                with tempfile.TemporaryDirectory(prefix="", dir=self.codegen.dist_dir) as tmp_dir:
-                    tmp_dir = Path(tmp_dir)
+                tmp_dir = tempfile.TemporaryDirectory(prefix="", dir=self.codegen.dist_dir)
+                tmp_dir = Path(tmp_dir.name)
+
+                archive_fname = self.codegen.write_archive(self.template_name, archive_format, tmp_dir)
+                # this is where streamlit serves static files
+                # ~/site-packages/streamlit/static/static/
+                dist_path = Path(st.__path__[0]) / "static/static" / tmp_dir
+
+                if not dist_path.is_dir():
+                    dist_path.mkdir(parents=True, exist_ok=True)
+
+                shutil.copy(archive_fname, dist_path)
+                st.success(f"Download link : [{archive_fname}](./static/{archive_fname})")
+
+                with col2:
+                    self.render_directory(Path(tmp_dir, self.template_name))
+
+    def download_panel(self):
+        st.markdown("")
+        archive_format = None
+        mimetype = ""
+
+        _, zip_col, tar_col, _ = st.beta_columns(4)
+        if zip_col.button("📦 Download zip"):
+            archive_format = "zip"
+            mimetype = "application/zip"
+
+        if tar_col.button("📦 Download tar"):
+            archive_format = "tar.gz"
+            mimetype = "application/x-tar"
 
-                    archive_fname = self.codegen.make_archive(self.template_name, archive_format, tmp_dir)
-                    # this is where streamlit serves static files
-                    # ~/site-packages/streamlit/static/static/
-                    dist_path = Path(st.__path__[0]) / "static/static" / tmp_dir
+        if archive_format is not None:
+            archive = self.codegen.writes_archive(self.template_name, archive_format)
 
-                    if not dist_path.is_dir():
-                        dist_path.mkdir(parents=True, exist_ok=True)
+            download_link = (
+                f'Download link: <a href="data:{mimetype};base64,{archive}" '
+                f'download="{self.template_name}.{archive_format}">'
+                f"{self.template_name}.{archive_format}</a>"
+            )
 
-                    shutil.copy(archive_fname, dist_path)
-                    st.success(f"Download link : [{archive_fname}](./static/{archive_fname})")
+            st.markdown(download_link, unsafe_allow_html=True)
 
-                    with col2:
-                        self.render_directory(Path(tmp_dir, self.template_name))
+    def add_download_panel(self):
+        if DEV_MODE:
+            self.download_panel_dev()
+        else:
+            self.download_panel()
 
     def run(self):
         self.add_sidebar()
         self.add_content()
-        self.add_download()
+        self.add_download_panel()
         st.info(TIP)