diff --git a/pystencils/backends/json.py b/pystencils/backends/json.py index ac2b2d4a055adb75e6d93272f25e6673e0b8ea6a..a0fcc9aa5f64109813429a3306c59385e5af47cd 100644 --- a/pystencils/backends/json.py +++ b/pystencils/backends/json.py @@ -9,6 +9,7 @@ """ import json +from pystencils.astnodes import NodeOrExpr from pystencils.backends.cbackend import CustomSympyPrinter, generate_c try: @@ -28,11 +29,19 @@ except Exception: def dumps(self, *args): raise ImportError('pyyaml not installed') - def dump(self, *args): - raise ImportError('pyyaml not installed') +def expr_to_dict(expr_or_node: NodeOrExpr, with_c_code=True, full_class_names=False): + """Converts a SymPy expression to a serializable dict (mainly for debugging purposes) + + The dict recursively contains all args of the expression as ``dict``s + + See :func:`.write_json` -def expr_to_dict(expr_or_node, with_c_code=True, full_class_names=False): + Args: + expr_or_node (NodeOrExpr): a SymPy expression or a :class:`pystencils.astnodes.Node` + with_c_code (bool, optional): include C representation of the nodes + full_class_names (bool, optional): use full class names (type(object) instead of ``type(object).__name__`` + """ self = {'str': str(expr_or_node)} if with_c_code: @@ -49,12 +58,26 @@ def expr_to_dict(expr_or_node, with_c_code=True, full_class_names=False): return self -def print_json(expr_or_node): +def print_json(expr_or_node: NodeOrExpr): + """Print debug JSON of an AST to string + + Args: + expr_or_node (NodeOrExpr): a SymPy expression or a :class:`pystencils.astnodes.Node` + + Returns: + str: JSON representation of AST + """ dict = expr_to_dict(expr_or_node) return json.dumps(dict, indent=4) -def write_json(filename, expr_or_node): +def write_json(filename: str, expr_or_node: NodeOrExpr): + """Writes debug JSON represenation of AST to file + + Args: + filename (str): Path for the file to write + expr_or_node (NodeOrExpr): a SymPy expression or a :class:`pystencils.astnodes.Node` + """ dict = expr_to_dict(expr_or_node) with open(filename, 'w') as f: json.dump(dict, f, indent=4)