Implement gzipping for static file serving

mesop/server/BUILD

@@ -24,3 +24,9 @@ py_test(
     srcs = ["wsgi_app_test.py"],
     deps = [":server"] + THIRD_PARTY_PY_PYTEST,
 )
+
+py_test(
+    name = "static_file_serving_test",
+    srcs = ["static_file_serving_test.py"],
+    deps = [":server"] + THIRD_PARTY_PY_PYTEST,
+)

mesop/server/static_file_serving.py

@@ -1,6 +1,8 @@
+import gzip
 import io
 import os
-from typing import Callable
+from io import BytesIO
+from typing import Any, Callable
 
 from flask import Flask, send_file
 from werkzeug.security import safe_join
@@ -49,11 +51,37 @@ def serve_root():
   def serve_file(path: str):
     preprocess_request()
     if is_file_path(path):
-      return send_file(get_path(path))
+      return send_file_compressed(get_path(path))
     else:
       return send_file(retrieve_index_html(), download_name="index.html")
 
 
 def is_file_path(path: str) -> bool:
   _, last_segment = os.path.split(path)
   return "." in last_segment
+
+
+# To avoid paying the cost of gzipping the same file multiple times
+# we have a singleton cache from file path to gzipped bytes.
+gzip_cache: dict[str, bytes] = {}
+
+
+def send_file_compressed(path: str) -> Any:
+  response = send_file(path)
+  response.headers["Content-Encoding"] = "gzip"
+  response.direct_passthrough = False
+  if path in gzip_cache:
+    gzip_data = gzip_cache[path]
+  else:
+    gzip_buffer = BytesIO()
+    with gzip.GzipFile(
+      mode="wb", fileobj=gzip_buffer, compresslevel=6
+    ) as gzip_file:
+      gzip_file.write(response.get_data())
+    gzip_buffer.seek(0)
+    gzip_data = gzip_buffer.getvalue()
+    gzip_cache[path] = gzip_data
+
+  response.set_data(gzip_data)
+  response.headers["Content-Length"] = str(len(response.get_data()))
+  return response

mesop/server/static_file_serving_test.py

@@ -0,0 +1,64 @@
+import gzip
+import tempfile
+from io import BytesIO
+
+from flask import Flask
+
+from mesop.server.static_file_serving import gzip_cache, send_file_compressed
+
+
+def test_send_file_compressed_uncached_request():
+  app = Flask(__name__)
+  with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+    file_content = b"abc"
+    tmp_file.write(file_content)
+    tmp_file_path = tmp_file.name
+
+  with app.test_request_context():
+    response = send_file_compressed(tmp_file_path)
+
+    assert response.headers["Content-Encoding"] == "gzip"
+    assert response.direct_passthrough is False
+    assert int(response.headers["Content-Length"]) == len(response.get_data())
+
+    # Check if the ungzipped data is correct
+    with gzip.GzipFile(fileobj=BytesIO(response.get_data()), mode="rb") as f:
+      ungzipped_data = f.read()
+      assert ungzipped_data == file_content
+
+    # Check that cache is properly stored
+    assert len(gzip_cache) == 1
+    assert gzip_cache[tmp_file_path] == response.get_data()
+
+
+def test_send_file_compressed_cached_request():
+  app = Flask(__name__)
+  with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+    tmp_file.write(b"")
+    tmp_file_path = tmp_file.name
+
+  cached_bytes = b"test_data"
+  gzip_buffer = BytesIO()
+  with gzip.GzipFile(
+    mode="wb", fileobj=gzip_buffer, compresslevel=6
+  ) as gzip_file:
+    gzip_file.write(cached_bytes)
+  gzip_buffer.seek(0)
+  gzip_cache[tmp_file_path] = gzip_buffer.getvalue()
+  with app.test_request_context():
+    response = send_file_compressed(tmp_file_path)
+
+    assert response.headers["Content-Encoding"] == "gzip"
+    assert response.direct_passthrough is False
+    assert int(response.headers["Content-Length"]) == len(response.get_data())
+
+    # Check that the cached bytes is returned
+    with gzip.GzipFile(fileobj=BytesIO(response.get_data()), mode="rb") as f:
+      ungzipped_data = f.read()
+      assert ungzipped_data == cached_bytes
+
+
+if __name__ == "__main__":
+  import pytest
+
+  raise SystemExit(pytest.main([__file__]))