Document backends.json

pystencils/backends/json.py

@@ -9,30 +9,22 @@
 """
 import json
 
+from pystencils.astnodes import NodeOrExpr
 from pystencils.backends.cbackend import CustomSympyPrinter, generate_c
 
-try:
-    import toml
-except Exception:
-    class toml:
-        def dumps(self, *args):
-            raise ImportError('toml not installed')
 
-        def dump(self, *args):
-            raise ImportError('toml 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)
 
-try:
-    import yaml
-except Exception:
-    class yaml:
-        def dumps(self, *args):
-            raise ImportError('pyyaml not installed')
+    The dict recursively contains all args of the expression as ``dict``s
 
-        def dump(self, *args):
-            raise ImportError('pyyaml not installed')
+    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,34 +41,26 @@ def expr_to_dict(expr_or_node, with_c_code=True, full_class_names=False):
     return self
 
 
-def print_json(expr_or_node):
-    dict = expr_to_dict(expr_or_node)
-    return json.dumps(dict, indent=4)
-
-
-def write_json(filename, expr_or_node):
-    dict = expr_to_dict(expr_or_node)
-    with open(filename, 'w') as f:
-        json.dump(dict, f, indent=4)
-
-
-def print_toml(expr_or_node):
-    dict = expr_to_dict(expr_or_node, full_class_names=False)
-    return toml.dumps(dict)
+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`
 
-def write_toml(filename, expr_or_node):
+    Returns:
+        str: JSON representation of AST
+    """
     dict = expr_to_dict(expr_or_node)
-    with open(filename, 'w') as f:
-        toml.dump(dict, f)
-
+    return json.dumps(dict, indent=4)
 
-def print_yaml(expr_or_node):
-    dict = expr_to_dict(expr_or_node, full_class_names=False)
-    return yaml.dump(dict)
 
+def write_json(filename: str, expr_or_node: NodeOrExpr):
+    """Writes debug JSON represenation of AST to file
 
-def write_yaml(filename, expr_or_node):
+    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:
-        yaml.dump(dict, f)
+        json.dump(dict, f, indent=4)