From 3a4f5a32f73f7e4712ace6009fdba9be21c345c5 Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Tue, 26 Mar 2024 15:13:42 -0700 Subject: [PATCH 1/5] Add orjson matrix to report builder test --- tests/test_report_builder.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_report_builder.py b/tests/test_report_builder.py index 43f3c24b..1f62a373 100644 --- a/tests/test_report_builder.py +++ b/tests/test_report_builder.py @@ -13,6 +13,7 @@ from viztracer.report_builder import ReportBuilder from .base_tmpl import BaseTmpl +from .package_env import package_matrix class TestReportBuilder(BaseTmpl): @@ -28,6 +29,7 @@ def test_file(self): result2 = s.getvalue() self.assertEqual(result1, result2) + @package_matrix(["~orjson", "orjson"]) def test_minimize_memory(self): json_path = os.path.join(os.path.dirname(__file__), "data", "multithread.json") with open(json_path) as f: @@ -72,6 +74,7 @@ def test_invalid_json(self): with self.assertRaises(Exception): ReportBuilder([invalid_json_path], verbose=1) + @package_matrix(["~orjson", "orjson"]) @patch('sys.stdout', new_callable=io.StringIO) def test_invalid_json_file(self, mock_stdout): with tempfile.TemporaryDirectory() as tmpdir: @@ -84,6 +87,7 @@ def test_invalid_json_file(self, mock_stdout): rb.save(s) self.assertIn("Invalid json file", mock_stdout.getvalue()) + @package_matrix(["~orjson", "orjson"]) @patch('sys.stdout', new_callable=io.StringIO) def test_all_invalid_json(self, mock_stdout): with tempfile.TemporaryDirectory() as tmpdir: @@ -95,6 +99,7 @@ def test_all_invalid_json(self, mock_stdout): rb.save(s) self.assertEqual(str(context.exception), "No valid json files found") + @package_matrix(["~orjson", "orjson"]) def test_combine(self): with tempfile.TemporaryDirectory() as tmpdir: file_path1 = os.path.join(tmpdir, "result1.json") From 5438ec352712b77034291c976945aa4753fd5eee Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Tue, 26 Mar 2024 15:52:25 -0700 Subject: [PATCH 2/5] Modify the orjson test --- tests/test_report_builder.py | 34 ++++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/tests/test_report_builder.py b/tests/test_report_builder.py index 1f62a373..c07fbb38 100644 --- a/tests/test_report_builder.py +++ b/tests/test_report_builder.py @@ -7,12 +7,14 @@ import os import shutil import tempfile +import textwrap from unittest.mock import patch import viztracer from viztracer.report_builder import ReportBuilder from .base_tmpl import BaseTmpl +from .cmdline_tmpl import CmdlineTmpl from .package_env import package_matrix @@ -29,7 +31,6 @@ def test_file(self): result2 = s.getvalue() self.assertEqual(result1, result2) - @package_matrix(["~orjson", "orjson"]) def test_minimize_memory(self): json_path = os.path.join(os.path.dirname(__file__), "data", "multithread.json") with open(json_path) as f: @@ -74,7 +75,6 @@ def test_invalid_json(self): with self.assertRaises(Exception): ReportBuilder([invalid_json_path], verbose=1) - @package_matrix(["~orjson", "orjson"]) @patch('sys.stdout', new_callable=io.StringIO) def test_invalid_json_file(self, mock_stdout): with tempfile.TemporaryDirectory() as tmpdir: @@ -87,7 +87,6 @@ def test_invalid_json_file(self, mock_stdout): rb.save(s) self.assertIn("Invalid json file", mock_stdout.getvalue()) - @package_matrix(["~orjson", "orjson"]) @patch('sys.stdout', new_callable=io.StringIO) def test_all_invalid_json(self, mock_stdout): with tempfile.TemporaryDirectory() as tmpdir: @@ -99,7 +98,6 @@ def test_all_invalid_json(self, mock_stdout): rb.save(s) self.assertEqual(str(context.exception), "No valid json files found") - @package_matrix(["~orjson", "orjson"]) def test_combine(self): with tempfile.TemporaryDirectory() as tmpdir: file_path1 = os.path.join(tmpdir, "result1.json") @@ -119,3 +117,31 @@ def test_combine(self): rb.save(output_file=s) data = json.loads(s.getvalue()) self.assertTrue(data["viztracer_metadata"]["overflow"]) + + +class TestReportBuilderCmdline(CmdlineTmpl): + @package_matrix(["~orjson", "orjson"]) + def test_package_matrix(self): + """ + The module will be imported only once so flipping the package matrix will only + work when we start a new script + """ + + with tempfile.TemporaryDirectory() as tmpdir: + invalid_json_path = os.path.join(os.path.dirname(__file__), "data", "fib.py") + invalid_json_file = shutil.copy(invalid_json_path, os.path.join(tmpdir, "invalid.json")) + + script = textwrap.dedent(f""" + import io + from viztracer.report_builder import ReportBuilder + rb = ReportBuilder(['{invalid_json_file}'], verbose=1) + try: + with io.StringIO() as s: + rb.save(s) + except Exception as e: + assert str(e) == "No valid json files found" + else: + assert False + """) + + self.template(["python", "cmdline_test.py"], script=script, expected_output_file=None) From 5312b46a98fbb878be0fda69754afcd966706415 Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Tue, 26 Mar 2024 21:25:30 -0700 Subject: [PATCH 3/5] import orjson as json --- src/viztracer/report_builder.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/viztracer/report_builder.py b/src/viztracer/report_builder.py index 72460a2c..62fd5338 100644 --- a/src/viztracer/report_builder.py +++ b/src/viztracer/report_builder.py @@ -2,14 +2,13 @@ # For details: https://github.com/gaogaotiantian/viztracer/blob/master/NOTICE.txt try: - import orjson # type: ignore + import orjson as json # type: ignore except ImportError: import json import gzip import os import re -import sys from string import Template from typing import Any, Dict, List, Optional, Sequence, TextIO, Tuple, Union @@ -27,10 +26,7 @@ def get_json(data: Union[Dict, str]) -> Dict[str, Any]: with open(data, encoding="utf-8") as f: json_str = f.read() - if "orjson" in sys.modules: - return orjson.loads(json_str) - else: - return json.loads(json_str) + return json.loads(json_str) class ReportBuilder: @@ -161,18 +157,18 @@ def generate_report( tmpl = f.read() with open(os.path.join(os.path.dirname(__file__), "html/trace_viewer_full.html"), encoding="utf-8") as f: sub["trace_viewer_full"] = f.read() - if "orjson" in sys.modules: - sub["json_data"] = orjson.dumps(self.combined_json) \ - .decode("utf-8") \ - .replace("", "<\\/script>") + if json.__name__ == "orjson": + sub["json_data"] = json.dumps(self.combined_json) \ + .decode("utf-8") \ + .replace("", "<\\/script>") else: sub["json_data"] = json.dumps(self.combined_json) \ .replace("", "<\\/script>") output_file.write(Template(tmpl).substitute(sub)) elif output_format == "json": self.prepare_json(file_info=file_info) - if "orjson" in sys.modules: - output_file.write(orjson.dumps(self.combined_json).decode("utf-8")) + if json.__name__ == "orjson": + output_file.write(json.dumps(self.combined_json).decode("utf-8")) else: if self.minimize_memory: json.dump(self.combined_json, output_file) # type: ignore From 542c6d3d07247abe93ae1a01d31e02dc94b819dc Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Tue, 26 Mar 2024 21:42:39 -0700 Subject: [PATCH 4/5] normpath for Windows --- tests/test_report_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_report_builder.py b/tests/test_report_builder.py index c07fbb38..2db4974f 100644 --- a/tests/test_report_builder.py +++ b/tests/test_report_builder.py @@ -134,7 +134,7 @@ def test_package_matrix(self): script = textwrap.dedent(f""" import io from viztracer.report_builder import ReportBuilder - rb = ReportBuilder(['{invalid_json_file}'], verbose=1) + rb = ReportBuilder(['{os.path.normpath(invalid_json_file)}'], verbose=1) try: with io.StringIO() as s: rb.save(s) From 08c81c80947cf46e02aeb9ee089c47922035c52a Mon Sep 17 00:00:00 2001 From: Tian Gao Date: Tue, 26 Mar 2024 21:59:31 -0700 Subject: [PATCH 5/5] Maybe use repr? --- tests/test_report_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_report_builder.py b/tests/test_report_builder.py index 2db4974f..4afed792 100644 --- a/tests/test_report_builder.py +++ b/tests/test_report_builder.py @@ -134,7 +134,7 @@ def test_package_matrix(self): script = textwrap.dedent(f""" import io from viztracer.report_builder import ReportBuilder - rb = ReportBuilder(['{os.path.normpath(invalid_json_file)}'], verbose=1) + rb = ReportBuilder([{repr(invalid_json_file)}], verbose=1) try: with io.StringIO() as s: rb.save(s)