diff --git a/midas/cli/main.py b/midas/cli/main.py index 220cca9..e683897 100644 --- a/midas/cli/main.py +++ b/midas/cli/main.py @@ -3,7 +3,7 @@ import json import logging from dataclasses import dataclass from pathlib import Path -from typing import Optional, TextIO +from typing import Optional, TextIO, get_args import click @@ -13,7 +13,13 @@ from midas.ast.location import Location from midas.ast.printer import MidasAstPrinter, PythonAstPrinter from midas.checker.checker import Checker from midas.checker.diagnostic import Diagnostic -from midas.cli.highlighter import DiagnosticsHighlighter, Highlighter, MidasHighlighter, PythonHighlighter +from midas.checker.types import Type +from midas.cli.highlighter import ( + DiagnosticsHighlighter, + Highlighter, + MidasHighlighter, + PythonHighlighter, +) from midas.lexer.midas import MidasLexer from midas.lexer.token import Token, TokenType from midas.parser.midas import MidasParser @@ -46,7 +52,9 @@ def compile(highlight: Optional[TextIO], file: TextIO): print( json.dumps( UniversalJSONDumper.dump( - checker.global_env, [("Environment", "_children")] + checker.global_env, + [("Environment", "_children")], + lambda obj: isinstance(obj, get_args(Type)), ), indent=4, ) diff --git a/midas/utils.py b/midas/utils.py index 0c387e2..bf10b95 100644 --- a/midas/utils.py +++ b/midas/utils.py @@ -1,18 +1,27 @@ -from typing import Any, Optional +from typing import Any, Callable, Optional + +AllowRepeat = Callable[[object], bool] class UniversalJSONDumper: @classmethod def dump( - cls, obj: Any, include_keys: Optional[list[str | tuple[str, str]]] = None + cls, + obj: Any, + include_keys: Optional[list[str | tuple[str, str]]] = None, + allow_repeat: Optional[AllowRepeat] = None, ) -> Any: if include_keys is None: include_keys = [] - return cls._dump(obj, include_keys, []) + return cls._dump(obj, include_keys, allow_repeat, []) @classmethod def _dump( - cls, obj: Any, include_keys: list[str | tuple[str, str]], visited: list[Any] + cls, + obj: Any, + include_keys: list[str | tuple[str, str]], + allow_repeat: Optional[AllowRepeat], + visited: list[Any], ) -> Any: if obj in visited: return None @@ -20,17 +29,22 @@ class UniversalJSONDumper: case str() | int() | float() | None: return obj case list() | set() | tuple(): - return [cls._dump(child, include_keys, visited) for child in obj] + return [ + cls._dump(child, include_keys, allow_repeat, visited) + for child in obj + ] case dict(): return { - str(k): cls._dump(v, include_keys, visited) for k, v in obj.items() + str(k): cls._dump(v, include_keys, allow_repeat, visited) + for k, v in obj.items() } case object(): - visited.append(obj) + if allow_repeat is None or not allow_repeat(obj): + visited.append(obj) return { "_type": obj.__class__.__name__, } | { - k: cls._dump(v, include_keys, visited) + k: cls._dump(v, include_keys, allow_repeat, visited) for k, v in obj.__dict__.items() if not k.startswith("_") or k in include_keys