✨ feat: add export json

chenjunjie · chenjunjie · commit 721dc8930cb8 · 2023-08-31T14:14:02.000+08:00
diff --git a/opendigger_pycli/cli/commands/export_cmd.py b/opendigger_pycli/cli/commands/export_cmd.py
@@ -1,29 +1,53 @@
 import typing as t
+from pathlib import Path
 
 import click
 
+from opendigger_pycli.exporters import (
+    SURPPORTED_EXPORT_FORMAT_TYPE,
+    SURPPORTED_EXPORT_FORMATS,
+)
+from opendigger_pycli.results.export import ExportResult
 from opendigger_pycli.utils.decorators import processor
 
 from ..base import pass_environment
 
 if t.TYPE_CHECKING:
+    from opendigger_pycli.results.query import QueryResults
+
     from ..base import Environment
 
 
 @click.command("export", help="Export metrics")
 @click.option(
     "--format",
     "-f",
-    "format_name",
-    type=click.Choice(["csv", "json", "mhtml"]),
+    type=click.Choice(SURPPORTED_EXPORT_FORMATS),
     required=True,
+    help="Format to export",
+)
+@click.option(
+    "--save-dir",
+    "-s",
+    "save_dir",
+    type=click.Path(file_okay=False, resolve_path=True, path_type=Path),
+    required=True,
+    help="Directory to save indicators",
+)
+@click.option(
+    "--split/--no-split",
+    "is_split",
+    default=True,
+    help="Save indicators in separate files",
 )
-@click.option("--filename", "-o", "filename", type=str, required=True)
 @processor
 @pass_environment
 def export(
     env: "Environment",
-    format_name: t.Literal["csv", "json", "mhtml"],
-    filename: click.Path,
+    results: "QueryResults",
+    format: SURPPORTED_EXPORT_FORMAT_TYPE,
+    save_dir: Path,
+    is_split: bool,
 ):
-    pass
+    ExportResult(results, format, save_dir, is_split).export()
+    yield from results
diff --git a/opendigger_pycli/datatypes/indicators/base.py b/opendigger_pycli/datatypes/indicators/base.py
@@ -113,6 +113,10 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return f"\033[34m{self.name}\033[0m: {self.value}"
 
+    @property
+    def tuple(self) -> t.Tuple[str, float]:
+        return (self.name, self.value)
+
 
 class NameNameAndValue(NamedTuple):
     name0: str
@@ -125,6 +129,10 @@ def __str__(self) -> str:
     def __repr__(self) -> str:
         return f"\033[34m{self.name0}\033[0m>>\033[31m{self.name1}\033[0m: {self.value}"
 
+    @property
+    def tuple(self) -> t.Tuple[str, str, float]:
+        return (self.name0, self.name1, self.value)
+
 
 AvgDataType = BaseData[float]
 LevelDataType = BaseData[List[int]]
diff --git a/opendigger_pycli/exporters/__init__.py b/opendigger_pycli/exporters/__init__.py
@@ -0,0 +1,4 @@
+import typing as t
+
+SURPPORTED_EXPORT_FORMATS = ["echarts", "json"]
+SURPPORTED_EXPORT_FORMAT_TYPE = t.Literal["echarts", "json"]
diff --git a/opendigger_pycli/exporters/json_exporter.py b/opendigger_pycli/exporters/json_exporter.py
@@ -0,0 +1,97 @@
+import json
+import typing as t
+from dataclasses import asdict
+
+if t.TYPE_CHECKING:
+    from opendigger_pycli.datatypes import (
+        BaseData,
+        BaseNetworkData,
+        NameAndValue,
+        NonTrivalNetworkInciatorData,
+        NonTrivialIndicatorData,
+        TrivialIndicatorData,
+        TrivialNetworkIndicatorData,
+    )
+
+
+def export_trivial_network_to_json(
+    indicator_data: "TrivialNetworkIndicatorData",
+) -> str:
+    result = {}
+    warmup = getattr(indicator_data.value, "nodes")[0]
+    if hasattr(warmup, "tuple"):
+        nodes: t.List[t.Tuple[str, float]] = []
+        edges: t.List[t.Tuple[str, str, float]] = []
+        base_network_data = t.cast("TrivialNetworkIndicatorData", indicator_data).value
+        for node in base_network_data.nodes:
+            nodes.append(node.tuple)
+        for edge in base_network_data.edges:
+            edges.append(edge.tuple)
+        result = {"nodes": nodes, "edges": edges}
+    else:
+        result = asdict(indicator_data)
+
+    return json.dumps(result, indent=2, sort_keys=True)
+
+
+def export_non_trivial_indicator_to_json(
+    indicator_data: "NonTrivialIndicatorData",
+) -> str:
+    result: t.Dict[str, t.Dict[str, t.Any]] = {}
+    indicator_data = t.cast("NonTrivialIndicatorData", indicator_data)
+    for key, values in indicator_data.value.items():
+        result[key] = {}
+        for base_data in values:  # type: ignore
+            value = t.cast("BaseData", base_data)
+            sub_key = (
+                f"{value.year}-{value.month:02}-raw"
+                if value.is_raw
+                else f"{value.year}-{value.month:02}"
+            )
+            result[key][sub_key] = value.value  # type: ignore
+    return json.dumps(result, indent=2, sort_keys=True)
+
+
+def export_indicator_to_json(
+    indicator_data: t.Union[
+        "TrivialIndicatorData",
+        "NonTrivialIndicatorData",
+        "TrivialNetworkIndicatorData",
+        "NonTrivalNetworkInciatorData",
+    ],
+) -> str:
+    if hasattr(indicator_data.value, "nodes"):
+        indicator_data = t.cast("TrivialNetworkIndicatorData", indicator_data)
+        return export_trivial_network_to_json(indicator_data)
+
+    if isinstance(indicator_data.value, dict):
+        indicator_data = t.cast("NonTrivialIndicatorData", indicator_data)
+        return export_non_trivial_indicator_to_json(indicator_data)
+
+    values = t.cast("t.List", indicator_data.value)
+    warmup = values[0]
+    result: t.Dict[str, t.Any] = {}
+    if hasattr(warmup.value, "nodes") or warmup.value is None:
+        base_data_list = t.cast("t.List[BaseData[BaseNetworkData]]", values)
+        for base_data in base_data_list:
+            key = (
+                f"{base_data.year}-{base_data.month:02}-raw"
+                if base_data.is_raw
+                else f"{base_data.year}-{base_data.month:02}"
+            )
+            result[key] = asdict(base_data.value) if base_data.value else None
+    else:
+        values = t.cast("t.List[BaseData]", values)
+        for value in values:
+            key = (
+                f"{value.year}-{value.month:02}-raw"
+                if value.is_raw
+                else f"{value.year}-{value.month:02}"
+            )
+            if isinstance(value.value, list) and hasattr(value.value[0], "name"):
+                value.value = t.cast("t.List[NameAndValue]", value.value)
+                result[key] = [v.tuple for v in value.value]
+            else:
+                result[key] = value.value
+
+    return json.dumps(result, indent=2, sort_keys=True)
diff --git a/opendigger_pycli/results/display.py b/opendigger_pycli/results/display.py
@@ -129,12 +129,12 @@ def _handle_save_path(
             query_result = t.cast(RepoQueryResult, query_result)
             save_path = str(
                 self.save_path
-                / f"{query_result.org_name}-{query_result.repo_name}-{self.mode}.html"
+                / f"repo-{query_result.org_name}-{query_result.repo_name}-{self.mode}.html"
             )
         else:
             query_result = t.cast(UserQueryResult, query_result)
             save_path = str(
-                self.save_path / f"{query_result.username}-{self.mode}.html"
+                self.save_path / f"user-{query_result.username}-{self.mode}.html"
             )
         return save_path
 
diff --git a/opendigger_pycli/results/export.py b/opendigger_pycli/results/export.py
@@ -1,7 +1,124 @@
+import json
+import typing as t
+
 from opendigger_pycli.console import CONSOLE
+from opendigger_pycli.console.utils import print_failed_query
+from opendigger_pycli.exporters.json_exporter import export_indicator_to_json
+
+from .query import QueryResults, RepoQueryResult, UserQueryResult
 
-from .query import RepoQueryResult, UserQueryResult
+if t.TYPE_CHECKING:
+    from pathlib import Path
+
+    from opendigger_pycli.exporters import (
+        SURPPORTED_EXPORT_FORMAT_TYPE,
+        SURPPORTED_EXPORT_FORMATS,
+    )
 
 
 class ExportResult:
-    pass
+    query_results: "QueryResults"
+    format: "SURPPORTED_EXPORT_FORMAT_TYPE"
+    save_path: "Path"
+    is_split: bool
+
+    def __init__(
+        self,
+        query_results: "QueryResults",
+        format: "SURPPORTED_EXPORT_FORMAT_TYPE",
+        save_path: "Path",
+        is_split: bool,
+        **kwargs,
+    ) -> None:
+        self.query_results = query_results
+        self.format = format
+        self.save_path = save_path
+        self.is_split = is_split
+
+    def _query_result_to_json(
+        self, query_result: t.Union["RepoQueryResult", "UserQueryResult"]
+    ) -> t.Dict[str, str]:
+        result: t.Dict[str, str] = {}
+        queried_indicators_data = query_result.queried_data
+        failed_queries = query_result.failed_query
+        for (
+            indicator_name,
+            indicator_dataloder_result,
+        ) in queried_indicators_data.items():
+            indicator_name_formated = indicator_name.replace("_", " ").title()
+            if (
+                not indicator_dataloder_result.is_success
+                or not indicator_dataloder_result.data
+            ):
+                CONSOLE.print(
+                    f"[red]Failed to load {indicator_name_formated} indicator data: {indicator_dataloder_result.desc}"
+                )
+                continue
+
+            failed_query = failed_queries.get(indicator_name)
+            if failed_query:
+                if isinstance(failed_query, dict):
+                    for key, value in failed_query.items():
+                        if value is not None:
+                            continue
+                        print_failed_query(indicator_name + "." + key, value)
+                else:
+                    print_failed_query(indicator_name, failed_query)
+
+            result[indicator_name] = export_indicator_to_json(
+                indicator_dataloder_result.data
+            )
+
+        return result
+
+    def _handle_save_path(
+        self, query_result: t.Union["RepoQueryResult", "UserQueryResult"]
+    ) -> t.Optional["Path"]:
+        if self.save_path is None:
+            return None
+
+        self.save_path.mkdir(parents=True, exist_ok=True)
+        if query_result.__class__ is RepoQueryResult:
+            query_result = t.cast(RepoQueryResult, query_result)
+            base_path = (
+                self.save_path
+                / f"repo-{query_result.org_name}-{query_result.repo_name}"
+            )
+        else:
+            query_result = t.cast(UserQueryResult, query_result)
+            base_path = self.save_path / f"user-{query_result.username}"
+
+        if self.is_split:
+            save_path = base_path / self.format
+            save_path.mkdir(parents=True, exist_ok=True)
+        else:
+            if self.format == "json":
+                save_path = base_path.with_suffix(".json")
+            else:
+                save_path = base_path.with_suffix(".html")
+
+        return save_path
+
+    def export(self) -> None:
+        if not self.query_results:
+            CONSOLE.print("[red]No results to export")
+            return
+
+        for query_result in self.query_results:
+            if self.format == "json":
+                result = self._query_result_to_json(query_result)
+                save_path = self._handle_save_path(query_result)
+                if save_path is None:
+                    raise ValueError("Save path is None")
+                if self.is_split:
+                    for indicator_name, indicator_json_data in result.items():
+                        save_path_splited = save_path / f"{indicator_name}.json"
+                        save_path_splited.write_text(indicator_json_data)
+                        CONSOLE.print(
+                            f"[green]Save Indicator {indicator_name} Data to {save_path}"
+                        )
+                else:
+                    save_path.write_text(json.dumps(result, indent=2, sort_keys=True))
+                    CONSOLE.print(f"[green]Save All Indicator Data to {save_path}")
+            else:
+                raise NotImplementedError

Original file line number	Diff line number	Diff line change
`@@ -129,12 +129,12 @@ def _handle_save_path(`
`129`	`129`	`query_result = t.cast(RepoQueryResult, query_result)`
`130`	`130`	`save_path = str(`
`131`	`131`	`self.save_path`
`132`		`- / f"{query_result.org_name}-{query_result.repo_name}-{self.mode}.html"`
	`132`	`+ / f"repo-{query_result.org_name}-{query_result.repo_name}-{self.mode}.html"`
`133`	`133`	`)`
`134`	`134`	`else:`
`135`	`135`	`query_result = t.cast(UserQueryResult, query_result)`
`136`	`136`	`save_path = str(`
`137`		`- self.save_path / f"{query_result.username}-{self.mode}.html"`
	`137`	`+ self.save_path / f"user-{query_result.username}-{self.mode}.html"`
`138`	`138`	`)`
`139`	`139`	`return save_path`
`140`	`140`