diff --git a/src/viztracer/viewer.py b/src/viztracer/viewer.py index 104d1739..734ea294 100644 --- 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 index 9af7acd1..f3e168e8 100644 --- 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 = ''