Support gzip for vizviewer (#544)

gaogaotiantian · Jan 26, 2025 · 6298bdc · 6298bdc
1 parent f745735
commit 6298bdc
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/src/viztracer/viewer.py b/src/viztracer/viewer.py
@@ -7,6 +7,7 @@
 import functools
 import html
 import http.server
+import gzip
 import io
 import json
 import os
@@ -301,15 +302,19 @@ def view(self) -> int:
         filename = os.path.basename(self.path)
 
         Handler: Callable[..., HttpHandler]
-        if filename.endswith("json"):
+        if filename.endswith("json") or filename.endswith("gz"):
             trace_data = None
             if self.use_external_procesor:
                 Handler = functools.partial(ExternalProcessorHandler, self)
                 self.externel_processor_process = ExternalProcessorProcess(self.path)
             else:
-                with open(self.path, encoding="utf-8", errors="ignore") as f:
-                    trace_data = json.load(f)
-                    self.file_info = trace_data.get("file_info", {})
+                if filename.endswith("gz"):
+                    with gzip.open(self.path, "rt", encoding="utf-8", errors="ignore") as f:
+                        trace_data = json.load(f)
+                else:
+                    with open(self.path, encoding="utf-8", errors="ignore") as f:
+                        trace_data = json.load(f)
+                self.file_info = trace_data.get("file_info", {})
                 Handler = functools.partial(PerfettoHandler, self)
         elif filename.endswith("html"):
             Handler = functools.partial(HtmlHandler, self)

diff --git a/tests/test_viewer.py b/tests/test_viewer.py
@@ -2,6 +2,7 @@
 # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt
 
 
+import gzip
 import json
 import multiprocessing
 import os
@@ -197,6 +198,27 @@ def test_json(self):
         finally:
             os.remove(f.name)
 
+    @unittest.skipIf(sys.platform == "win32", "Can't send Ctrl+C reliably on Windows")
+    def test_gz(self):
+        json_script = '{"file_info": {}, "traceEvents": []}'
+        with tempfile.TemporaryDirectory() as tmpdir:
+            filename = os.path.join(tmpdir, "test.json.gz")
+            with gzip.open(filename, "wt") as f:
+                f.write(json_script)
+            v = Viewer(filename)
+            try:
+                v.run()
+                time.sleep(0.5)
+                resp = urllib.request.urlopen(v.url())
+                self.assertTrue(resp.code == 200)
+                resp = urllib.request.urlopen(f"{v.url()}/file_info")
+                self.assertEqual(json.loads(resp.read().decode("utf-8")), {})
+                resp = urllib.request.urlopen(f"{v.url()}/localtrace")
+                self.assertEqual(json.loads(gzip.decompress(resp.read()).decode("utf-8")),
+                                 json.loads(json_script))
+            finally:
+                v.stop()
+
     @unittest.skipIf(sys.platform == "win32", "Can't send Ctrl+C reliably on Windows")
     def test_html(self):
         html = '<html></html>'