feat flag to include/exclude functions in feature model

src/feature_model.py

+# Configuration
+model_name = 'SimpleHeatingSystem.mo'
+output_file = 'feature_model.json'
+include_equations = True  # set to "False" if equations should be excluded
+
 from anytree import Node, RenderTree
 import json
 from parse_modelica import parse_model
@@ -15,6 +20,10 @@ class FeatureModel:
         for param_name, param_value in parameters.items():
             Node(f"{param_name}={param_value}", parent=component_node)
 
+    # Method to add an equation to the tree
+    def add_equation(self, equation, index):
+        Node(f"Equation {index+1}: {equation}", parent=self.root)
+
     # Method to display the tree structure in the console
     def display(self):
         for pre, fill, node in RenderTree(self.root):
@@ -34,13 +43,17 @@ class FeatureModel:
 
 if __name__ == "__main__":
     # Parse the Modelica file and extract components
-    components = parse_model('SimpleHeatingSystem.mo')
+    components, equations = parse_model(model_name, include_equations)
 
     # Create a feature model and add the extracted components
-    feature_model = FeatureModel("SimpleHeatingSystem")
+    feature_model = FeatureModel(model_name.split('.')[0])
     for component_type, component_name, parameters in components:
         feature_model.add_component(component_type, component_name, parameters)
 
+    if include_equations:
+        for index, equation in enumerate(equations):
+            feature_model.add_equation(equation, index)
+
     # Display and save the feature tree
     feature_model.display()
-    feature_model.save_to_json('feature_model.json')
\ No newline at end of file
+    feature_model.save_to_json(output_file)
\ No newline at end of file

src/parse_modelica.py

 # This parses a Modelica model to extract its components (variables and parameters) and their parameters. 
 # It uses the ANTLR-generated lexer and parser to build a parse tree and then walks the tree to extract relevant information.
 
+# Configuration
+file_name = 'SimpleHeatingSystem.mo'
+
 import sys
 import os
 from antlr4 import *
@@ -12,15 +15,17 @@ from modelicaListener import modelicaListener
 
 # Listener class to extract components and their parameters from the parse tree
 class ModelicaFeatureExtractor(modelicaListener):
-    def __init__(self):
+    def __init__(self, include_equations):
         self.components = []  # List to store extracted components
+        self.equations = []  # List to store extracted equations
+        self.include_equations = include_equations
 
     # Method called when entering a component clause in the parse tree
     def enterComponent_clause(self, ctx: modelicaParser.Component_clauseContext):
         # Extract the component type and name
         type_prefix = ctx.type_prefix().getText() if ctx.type_prefix() else ""
         component_type = type_prefix + " " + ctx.type_specifier().getText()
-        print(f"Component type: {component_type}")
+        # print(f"Component type: {component_type}")
 
         # Iterate over each component declaration in the component list
         for component in ctx.component_list().component_declaration():
@@ -34,11 +39,11 @@ class ModelicaFeatureExtractor(modelicaListener):
                 modification = component.declaration().modification()
                 param_value = self.get_modification_value(modification)
                 if param_value is not None:
-                    parameters[component_name] = param_value  # Store the parameter value
+                    parameters["value"] = param_value  # Store the parameter value
                     # print(f"Param: {component_name} = {param_value}")  uncomment for debugging
 
             # Add the component to the list
-            self.components.append((component_type, component_name, parameters))
+            self.components.append((component_type.strip(), component_name, parameters))
 
     # Helper method to extract the value from a modification context
     def get_modification_value(self, modification_ctx):
@@ -59,8 +64,15 @@ class ModelicaFeatureExtractor(modelicaListener):
                 return "({})".format(", ".join(values))
         return None
 
+    # Method to extract equations
+    def enterEquation_section(self, ctx: modelicaParser.Equation_sectionContext):
+        if self.include_equations:
+            for equation in ctx.equation():
+                equations_text = equation.getText()
+                self.equations.append(equations_text)
+
 # Function to parse a Modelica file and extract components
-def parse_model(file_name):
+def parse_model(file_name, include_equations):
 
     input_stream = FileStream(file_name, encoding='utf-8')  # Read the file
     lexer = modelicaLexer(input_stream)  # Tokenize the input
@@ -68,14 +80,15 @@ def parse_model(file_name):
     parser = modelicaParser(stream)  # Create a parser
     tree = parser.stored_definition()  # Parse the input into a parse tree
 
-    extractor = ModelicaFeatureExtractor()  # Create a listener
+    extractor = ModelicaFeatureExtractor(include_equations)  # Create a listener
     walker = ParseTreeWalker()  # Create a parse tree walker
     walker.walk(extractor, tree)  # Walk the parse tree with the listener
 
-    return extractor.components  # Return the extracted components
+    return extractor.components, extractor.equations  # Return the extracted components
 
 if __name__ == "__main__":
-    file_name = 'SimpleHeatingSystem.mo'
-    components = parse_model(file_name)  # Parse the file and extract components
+    components, equations = parse_model(file_name, include_equations)  # Parse the file and extract components
     for component in components:
         print(component)  # Print the extracted components
+    for equation in equations:
+        print(equation)
\ No newline at end of file