Coverage for skema/program_analysis/CAST/matlab/variable_context.py: 83%

1from typing import List, Dict, Set 

2from skema.program_analysis.CAST2FN.model.cast import ( 

3 Var, 

4 Name, 

5) 

6 

7class VariableContext(object): 

8 def __init__(self): 

9 self.context = [{}] # Stack of context dictionaries 

10 self.context_return_values = [set()] # Stack of context return values 

11 self.all_symbols = {} 

12 self.record_definitions = {} 

13 

14 # The prefix is used to handle adding Record types to the variable context. 

15 # This gives each symbol a unqique name. For example "a" would become "type_name.a" 

16 # For nested type definitions (derived type in a module), multiple prefixes can be added. 

17 self.prefix = [] 

18 

19 # Flag neccessary to declare if a function is internal or external 

20 self.internal = False 

21 

22 self.variable_id = 0 

23 self.iterator_id = 0 

24 self.stop_condition_id = 0 

25 

26 def push_context(self): 

27 """Create a new variable context and add it to the stack""" 

28 

29 self.context.append({}) 

30 self.context_return_values.append(set()) 

31 

32 def pop_context(self): 

33 """Pop the current variable context off of the stack and remove any references to those symbols.""" 

34 

35 context = self.context.pop() 

36 

37 # Remove symbols from all_symbols variable 

38 for symbol in context: 

39 self.all_symbols.pop(symbol) 

40 

41 self.context_return_values.pop() 

42 

43 def add_variable(self, symbol: str, type: str, source_refs: List) -> Name: 

44 """Add a variable to the current variable context""" 

45 # Generate the full symbol name using the prefix 

46 full_symbol_name = ".".join(self.prefix + [symbol]) 

47 

48 cast_name = Name(source_refs=source_refs) 

49 cast_name.name = symbol 

50 cast_name.id = self.variable_id 

51 

52 # Update variable id 

53 self.variable_id += 1 

54 

55 # Add the node to the variable context 

56 self.context[-1][full_symbol_name] = { 

57 "node": cast_name, 

58 "type": type, 

59 } 

60 

61 # Add reference to all_symbols 

62 self.all_symbols[full_symbol_name] = self.context[-1][full_symbol_name] 

63 

64 return cast_name 

65 

66 def is_variable(self, symbol: str) -> bool: 

67 """Check if a symbol exists in any context""" 

68 return symbol in self.all_symbols 

69 

70 def get_node(self, symbol: str) -> Dict: 

71 return self.all_symbols[symbol]["node"] 

72 

73 def get_type(self, symbol: str) -> str: 

74 return self.all_symbols[symbol]["type"] 

75 

76 def get_gromet_function_node(self, func_name: str) -> Name: 

77 if self.is_variable(func_name): 

78 return self.get_node(func_name) 

79 

80 def generate_iterator(self): 

81 symbol = f"generated_iter_{self.iterator_id}" 

82 self.iterator_id += 1 

83 return self.add_variable(symbol, "iterator", None) 

84 

85 def generate_stop_condition(self): 

86 symbol = f"sc_{self.stop_condition_id}" 

87 self.stop_condition_id += 1 

88 return self.add_variable(symbol, "boolean", None) 

89