refactor: Introduce PythonCodeGeneratorContext to manage per-namespace generation state and extract common constants.

Author: dschwartz-ftadvisory

src/main/java/com/regnosys/rosetta/generator/python/PythonCodeGenerator.java

@@ -1,7 +1,6 @@
 package com.regnosys.rosetta.generator.python;
-// TODO: collect imports as a set rather than an array
-
 // TODO: re-engineer type generation to use an object that has the features carried throughout the generation (imports, etc.)
+
 // TODO: function support
 // TODO: review and consolidate unit tests
 // TODO: review migrating choice alias processor to PythonModelObjectGenerator
@@ -13,6 +12,8 @@
 import com.regnosys.rosetta.generator.python.object.PythonModelObjectGenerator;
 import com.regnosys.rosetta.generator.python.util.PythonCodeGeneratorUtil;
 import com.regnosys.rosetta.generator.python.util.PythonCodeWriter;
+import static com.regnosys.rosetta.generator.python.util.PythonCodeGeneratorConstants.*;
+
 import com.regnosys.rosetta.rosetta.RosettaEnumeration;
 import com.regnosys.rosetta.rosetta.RosettaModel;
 import com.regnosys.rosetta.rosetta.simple.Data;
@@ -25,7 +26,7 @@
 
 import org.jgrapht.Graph;
 import org.jgrapht.graph.DefaultEdge;
-import org.jgrapht.graph.DirectedAcyclicGraph;
+
 import org.jgrapht.traverse.TopologicalOrderIterator;
 
 import java.util.*;
@@ -95,29 +96,35 @@ public class PythonCodeGenerator extends AbstractExternalGenerator {
 
     private static final Logger LOGGER = LoggerFactory.getLogger(PythonCodeGenerator.class);
 
-    private List<String> subfolders = null;
-    private Map<String, Map<String, CharSequence>> objects = null; // Python code for types by nameSpace, by type name
-    private Graph<String, DefaultEdge> dependencyDAG = null;
-    private Set<String> enumImports = null;
+    private Map<String, PythonCodeGeneratorContext> contexts = null;
 
     public PythonCodeGenerator() {
-        super("python");
+        super(PYTHON);
     }
 
     @Override
     public Map<String, ? extends CharSequence> beforeAllGenerate(ResourceSet set,
             Collection<? extends RosettaModel> models, String version) {
-        subfolders = new ArrayList<>();
-        objects = new HashMap<>();
-        dependencyDAG = new DirectedAcyclicGraph<>(DefaultEdge.class);
-        enumImports = new HashSet<>();
-        pojoGenerator.beforeAllGenerate(dependencyDAG, enumImports);
+
+        contexts = new HashMap<>();
         return Collections.emptyMap();
     }
 
     @Override
     public Map<String, ? extends CharSequence> generate(Resource resource, RosettaModel model, String version) {
-        String cleanVersion = cleanVersion(version);
+        if (model == null) {
+            throw new IllegalArgumentException("Model is null");
+        }
+        LOGGER.debug("Processing module: {}", model.getName());
+
+        String nameSpace = PythonCodeGeneratorUtil.getNamespace(model);
+        PythonCodeGeneratorContext context = contexts.get(nameSpace);
+        if (context == null) {
+            context = new PythonCodeGeneratorContext();
+            contexts.put(nameSpace, context);
+        }
+
+        String cleanVersion = PythonCodeGeneratorUtil.cleanVersion(version);
 
         Map<String, CharSequence> result = new HashMap<>();
 
@@ -132,21 +139,15 @@ public PythonCodeGenerator() {
                 .map(Function.class::cast).collect(Collectors.toList());
 
         if (!rosettaClasses.isEmpty() || !rosettaEnums.isEmpty() || !rosettaFunctions.isEmpty()) {
-            addSubfolder(model.getName());
+            context.addSubfolder(model.getName());
             if (!rosettaFunctions.isEmpty()) {
-                addSubfolder(model.getName() + ".functions");
+                context.addSubfolder(model.getName() + ".functions");
             }
         }
 
-        LOGGER.debug("Processing module: {}", model.getName());
-
-        String nameSpace = PythonCodeGeneratorUtil.getNamespace(model);
-        Map<String, CharSequence> currentObject = objects.get(nameSpace);
-        if (currentObject == null) {
-            currentObject = new HashMap<String, CharSequence>();
-            objects.put(nameSpace, currentObject);
-        }
-        currentObject.putAll(pojoGenerator.generate(rosettaClasses, cleanVersion));
+        Map<String, CharSequence> currentObject = context.getObjects();
+        currentObject.putAll(pojoGenerator.generate(rosettaClasses, cleanVersion, context.getDependencyDAG(),
+                context.getEnumImports()));
         result.putAll(enumGenerator.generate(rosettaEnums, cleanVersion));
         result.putAll(funcGenerator.generate(rosettaFunctions, cleanVersion));
 
@@ -158,26 +159,30 @@ public PythonCodeGenerator() {
             ResourceSet set,
             Collection<? extends RosettaModel> models,
             String version) {
-        String cleanVersion = cleanVersion(version);
         Map<String, CharSequence> result = new HashMap<>();
-
-        List<String> workspaces = getWorkspaces(subfolders);
-        result.putAll(generateWorkspaces(workspaces, cleanVersion));
-        result.putAll(generateInits(subfolders));
-
-        for (String nameSpace : objects.keySet()) {
-            Map<String, CharSequence> currentObject = objects.get(nameSpace);
-            if (currentObject != null && !currentObject.isEmpty()) {
-                result.put("pyproject.toml", PythonCodeGeneratorUtil.createPYProjectTomlFile(nameSpace, cleanVersion));
-                result.putAll(processDAG(nameSpace, currentObject));
-            }
+        String cleanVersion = PythonCodeGeneratorUtil.cleanVersion(version);
+        for (String nameSpace : contexts.keySet()) {
+            PythonCodeGeneratorContext context = contexts.get(nameSpace);
+            List<String> subfolders = context.getSubfolders();
+            result.putAll(generateWorkspaces(getWorkspaces(subfolders), cleanVersion));
+            result.putAll(generateInits(subfolders));
+            result.putAll(processDAG(nameSpace, context, cleanVersion));
         }
         return result;
     }
 
-    private Map<String, CharSequence> processDAG(String nameSpace, Map<String, CharSequence> nameSpaceObjects) {
+    private Map<String, CharSequence> processDAG(String nameSpace, PythonCodeGeneratorContext context,
+            String cleanVersion) {
+        if (nameSpace == null || context == null || cleanVersion == null) {
+            throw new IllegalArgumentException("Invalid arguments");
+        }
         Map<String, CharSequence> result = new HashMap<>();
-        if (dependencyDAG != null) {
+        Map<String, CharSequence> nameSpaceObjects = context.getObjects();
+        Graph<String, DefaultEdge> dependencyDAG = context.getDependencyDAG();
+        Set<String> enumImports = context.getEnumImports();
+
+        if (nameSpaceObjects != null && !nameSpaceObjects.isEmpty() && dependencyDAG != null && enumImports != null) {
+            result.put(PYPROJECT_TOML, PythonCodeGeneratorUtil.createPYProjectTomlFile(nameSpace, cleanVersion));
             PythonCodeWriter bundleWriter = new PythonCodeWriter();
             TopologicalOrderIterator<String, DefaultEdge> topologicalOrderIterator = new TopologicalOrderIterator<>(
                     dependencyDAG);
@@ -205,7 +210,7 @@ private Map<String, CharSequence> processDAG(String nameSpace, Map<String, CharS
 
                     // create the stub
                     String[] parsedName = name.split("\\.");
-                    String stubFileName = "src/" + String.join("/", parsedName) + ".py";
+                    String stubFileName = SRC + String.join("/", parsedName) + ".py";
 
                     PythonCodeWriter stubWriter = new PythonCodeWriter();
                     stubWriter.appendLine("# pylint: disable=unused-import");
@@ -224,25 +229,11 @@ private Map<String, CharSequence> processDAG(String nameSpace, Map<String, CharS
             }
             bundleWriter.newLine();
             bundleWriter.appendLine("# EOF");
-            result.put("src/" + nameSpace + "/_bundle.py", bundleWriter.toString());
+            result.put(SRC + nameSpace + "/_bundle.py", bundleWriter.toString());
         }
         return result;
     }
 
-    private String cleanVersion(String version) {
-        if (version == null || version.equals("${project.version}")) {
-            return "0.0.0";
-        }
-
-        String[] versionParts = version.split("\\.");
-        if (versionParts.length > 2) {
-            String thirdPart = versionParts[2].replaceAll("[^\\d]", "");
-            return versionParts[0] + "." + versionParts[1] + "." + thirdPart;
-        }
-
-        return "0.0.0";
-    }
-
     private List<String> getWorkspaces(List<String> subfolders) {
         return subfolders.stream().map(subfolder -> subfolder.split("\\.")[0]).distinct().collect(Collectors.toList());
     }
@@ -251,7 +242,7 @@ private Map<String, String> generateWorkspaces(List<String> workspaces, String v
         Map<String, String> result = new HashMap<>();
 
         for (String workspace : workspaces) {
-            result.put(PythonCodeGeneratorUtil.toPyFileName(workspace, "__init__"),
+            result.put(PythonCodeGeneratorUtil.toPyFileName(workspace, INIT),
                     PythonCodeGeneratorUtil.createTopLevelInitFile(version));
             result.put(PythonCodeGeneratorUtil.toPyFileName(workspace, "version"),
                     PythonCodeGeneratorUtil.createVersionFile(version));
@@ -268,16 +259,11 @@ private Map<String, String> generateInits(List<String> subfolders) {
             String[] parts = subfolder.split("\\.");
             for (int i = 1; i < parts.length; i++) {
                 String key = String.join(".", Arrays.copyOfRange(parts, 0, i + 1));
-                result.putIfAbsent(PythonCodeGeneratorUtil.toPyFileName(key, "__init__"), " ");
+                result.putIfAbsent(PythonCodeGeneratorUtil.toPyFileName(key, INIT), " ");
             }
         }
 
         return result;
     }
 
-    private void addSubfolder(String subfolder) {
-        if (!subfolders.contains(subfolder)) {
-            subfolders.add(subfolder);
-        }
-    }
 }

src/main/java/com/regnosys/rosetta/generator/python/PythonCodeGeneratorContext.java

@@ -0,0 +1,48 @@
+package com.regnosys.rosetta.generator.python;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import org.jgrapht.Graph;
+import org.jgrapht.graph.DefaultEdge;
+import org.jgrapht.graph.DirectedAcyclicGraph;
+
+public class PythonCodeGeneratorContext {
+    private List<String> subfolders = null;
+    private Map<String, CharSequence> objects = null; // Python code for types by nameSpace, by type name
+    private Graph<String, DefaultEdge> dependencyDAG = null;
+    private Set<String> enumImports = null;
+
+    public PythonCodeGeneratorContext() {
+        this.subfolders = new ArrayList<>();
+        this.objects = new HashMap<>();
+        this.dependencyDAG = new DirectedAcyclicGraph<>(DefaultEdge.class);
+        this.enumImports = new HashSet<>();
+    }
+
+    public List<String> getSubfolders() {
+        return subfolders;
+    }
+
+    public Map<String, CharSequence> getObjects() {
+        return objects;
+    }
+
+    public Graph<String, DefaultEdge> getDependencyDAG() {
+        return dependencyDAG;
+    }
+
+    public Set<String> getEnumImports() {
+        return enumImports;
+    }
+
+    public void addSubfolder(String subfolder) {
+        if (!subfolders.contains(subfolder)) {
+            subfolders.add(subfolder);
+        }
+    }
+}

src/main/java/com/regnosys/rosetta/generator/python/object/PythonModelObjectGenerator.java

@@ -10,9 +10,7 @@
 import jakarta.inject.Inject;
 import org.jgrapht.Graph;
 import org.jgrapht.graph.DefaultEdge;
-import org.jgrapht.graph.DirectedAcyclicGraph;
 import org.jgrapht.graph.GraphCycleProhibitedException;
-import org.jgrapht.traverse.TopologicalOrderIterator;
 
 import java.util.*;
 import java.util.stream.Collectors;
@@ -34,24 +32,15 @@ public class PythonModelObjectGenerator {
     @Inject
     private PythonChoiceAliasProcessor pythonChoiceAliasProcessor;
 
-    private Graph<String, DefaultEdge> dependencyDAG = null;
-    private Set<String> enumImports = null;
-
-    public void beforeAllGenerate(Graph<String, DefaultEdge> dependencyDAGIn, Set<String> enumImportsIn) {
-        dependencyDAG = dependencyDAGIn;
-        enumImports = enumImportsIn;
-    }
-
     /**
      * Generate Python from the collection of Rosetta classes (of type Data).
-     * Note: this function updates the dependency graph used by afterAllGenerate to
-     * create the bundle
      * 
      * @param rosettaClasses the collection of Rosetta Classes for this model
      * @param version        the version for this collection of classes
      * @return a Map of all the generated Python indexed by the class name
      */
-    public Map<String, String> generate(Iterable<Data> rosettaClasses, String version) {
+    public Map<String, String> generate(Iterable<Data> rosettaClasses, String version,
+            Graph<String, DefaultEdge> dependencyDAG, Set<String> enumImports) {
         if (dependencyDAG == null) {
             throw new RuntimeException("Dependency DAG not initialized");
         }
@@ -65,7 +54,7 @@ public Map<String, String> generate(Iterable<Data> rosettaClasses, String versio
             String nameSpace = PythonCodeGeneratorUtil.getNamespace(model);
 
             // Generate Python for the class
-            String pythonClass = generateClass(rosettaClass, nameSpace, version);
+            String pythonClass = generateClass(rosettaClass, nameSpace, version, enumImports);
 
             // construct the class name using "." as a delimiter
             String className = model.getName() + "." + rosettaClass.getName();
@@ -76,13 +65,14 @@ public Map<String, String> generate(Iterable<Data> rosettaClasses, String versio
                 Data superClass = rosettaClass.getSuperType();
                 RosettaModel superModel = (RosettaModel) superClass.eContainer();
                 String superClassName = superModel.getName() + "." + superClass.getName();
-                addDependency(className, superClassName);
+
+                addDependency(dependencyDAG, className, superClassName);
             }
         }
         return result;
     }
 
-    private void addDependency(String className, String dependencyName) {
+    private void addDependency(Graph<String, DefaultEdge> dependencyDAG, String className, String dependencyName) {
         dependencyDAG.addVertex(dependencyName);
         if (!className.equals(dependencyName)) {
             try {
@@ -93,11 +83,17 @@ private void addDependency(String className, String dependencyName) {
         }
     }
 
-    private String generateClass(Data rosettaClass, String nameSpace, String version) {
+    private String generateClass(Data rosettaClass, String nameSpace, String version, Set<String> enumImports) {
+        if (rosettaClass == null) {
+            throw new RuntimeException("Rosetta class not initialized");
+        }
         if (rosettaClass.getSuperType() != null && rosettaClass.getSuperType().getName() == null) {
             throw new RuntimeException(
                     "The class superType for " + rosettaClass.getName() + " exists but its name is null");
         }
+        if (enumImports == null) {
+            throw new RuntimeException("Enum imports not initialized");
+        }
 
         Set<String> enumImportsFound = pythonAttributeProcessor.getImportsFromAttributes(rosettaClass);
         enumImports.addAll(enumImportsFound);

src/main/java/com/regnosys/rosetta/generator/python/util/PythonCodeGeneratorConstants.java

@@ -0,0 +1,12 @@
+package com.regnosys.rosetta.generator.python.util;
+
+public final class PythonCodeGeneratorConstants {
+    private PythonCodeGeneratorConstants() {
+        // Restricted constructor
+    }
+
+    public static final String PYTHON = "python";
+    public static final String SRC = "src/";
+    public static final String PYPROJECT_TOML = "pyproject.toml";
+    public static final String INIT = "__init__";
+}

src/main/java/com/regnosys/rosetta/generator/python/util/PythonCodeGeneratorUtil.java

@@ -150,4 +150,18 @@ public static String createPYProjectTomlFile(String namespace, String version) {
                 [tool.setuptools.packages.find]
                 where = ["src"]""".formatted(namespace, version).stripIndent();
     }
+
+    public static String cleanVersion(String version) {
+        if (version == null || version.equals("${project.version}")) {
+            return "0.0.0";
+        }
+
+        String[] versionParts = version.split("\\.");
+        if (versionParts.length > 2) {
+            String thirdPart = versionParts[2].replaceAll("[^\\d]", "");
+            return versionParts[0] + "." + versionParts[1] + "." + thirdPart;
+        }
+
+        return "0.0.0";
+    }
 }