1# -*- coding: utf-8 -*- 

2# 

3# Copyright © 2019 Stephan Seitz <stephan.seitz@fau.de> 

4# 

5# Distributed under terms of the GPLv3 license. 

6 

7""" 

8 

9""" 

10import json 

11 

12from pystencils.astnodes import NodeOrExpr 

13from pystencils.backends.cbackend import CustomSympyPrinter, generate_c 

14 

15try: 

16 import yaml 

17except ImportError: 

18 raise ImportError('yaml not installed') 

19 

20 

21def expr_to_dict(expr_or_node: NodeOrExpr, with_c_code=True, full_class_names=False): 

22 """Converts a SymPy expression to a serializable dict (mainly for debugging purposes) 

23 

24 The dict recursively contains all args of the expression as ``dict``s 

25 

26 See :func:`.write_json` 

27 

28 Args: 

29 expr_or_node (NodeOrExpr): a SymPy expression or a :class:`pystencils.astnodes.Node` 

30 with_c_code (bool, optional): include C representation of the nodes 

31 full_class_names (bool, optional): use full class names (type(object) instead of ``type(object).__name__`` 

32 """ 

33 

34 self = {'str': str(expr_or_node)} 

35 if with_c_code: 

36 try: 

37 self.update({'c': generate_c(expr_or_node)}) 

38 except Exception: 

39 try: 

40 self.update({'c': CustomSympyPrinter().doprint(expr_or_node)}) 

41 except Exception: 

42 pass 

43 for a in expr_or_node.args: 

44 self.update({str(a.__class__ if full_class_names else a.__class__.__name__): expr_to_dict(a)}) 

45 

46 return self 

47 

48 

49def print_json(expr_or_node: NodeOrExpr): 

50 """Print debug JSON of an AST to string 

51 

52 Args: 

53 expr_or_node (NodeOrExpr): a SymPy expression or a :class:`pystencils.astnodes.Node` 

54 

55 Returns: 

56 str: JSON representation of AST 

57 """ 

58 expr_or_node_dict = expr_to_dict(expr_or_node) 

59 return json.dumps(expr_or_node_dict, indent=4) 

60 

61 

62def write_json(filename: str, expr_or_node: NodeOrExpr): 

63 """Writes debug JSON represenation of AST to file 

64 

65 Args: 

66 filename (str): Path for the file to write 

67 expr_or_node (NodeOrExpr): a SymPy expression or a :class:`pystencils.astnodes.Node` 

68 """ 

69 expr_or_node_dict = expr_to_dict(expr_or_node) 

70 with open(filename, 'w') as f: 

71 json.dump(expr_or_node_dict, f, indent=4) 

72 

73 

74def print_yaml(expr_or_node): 

75 expr_or_node_dict = expr_to_dict(expr_or_node, full_class_names=False) 

76 return yaml.dump(expr_or_node_dict) 

77 

78 

79def write_yaml(filename, expr_or_node): 

80 expr_or_node_dict = expr_to_dict(expr_or_node) 

81 with open(filename, 'w') as f: 

82 yaml.dump(expr_or_node_dict, f)