Merge pull request #6 from priyanka-amarnani/feature/dsl-bundle-upgrade-switch-fix

Fix for #5

Author: dschwartz-ftadvisory

src/main/java/com/regnosys/rosetta/generator/python/expressions/PythonExpressionGenerator.xtend

@@ -46,6 +46,7 @@ import com.regnosys.rosetta.rosetta.expression.MinOperation
 import com.regnosys.rosetta.rosetta.expression.MaxOperation
 import com.regnosys.rosetta.rosetta.expression.SwitchOperation
 import com.regnosys.rosetta.rosetta.expression.SwitchCaseGuard
+import com.regnosys.rosetta.rosetta.expression.SwitchCaseOrDefault
 import com.regnosys.rosetta.rosetta.simple.Attribute
 import com.regnosys.rosetta.rosetta.simple.Condition
 import com.regnosys.rosetta.rosetta.simple.Data
@@ -61,7 +62,7 @@ class PythonExpressionGenerator {
 
     public var List<String> importsFound
     public var ifCondBlocks = new ArrayList<String>()
-    public var switchCondBlocks = new ArrayList<String>()
+    public var isSwitchCond = false
 
     def String generateConditions(Data cls) {
         var nConditions = 0;
@@ -198,12 +199,10 @@ class PythonExpressionGenerator {
 
     private def generateIfThenElseOrSwitch(Condition c) {
         ifCondBlocks.clear()
-        switchCondBlocks.clear()
+        isSwitchCond=false
+        
         var expr = generateExpression(c.expression, 0, false)
-        if (!switchCondBlocks.isEmpty()) {
-            var switchBlocks = '''    «FOR arg : switchCondBlocks»«arg»«ENDFOR»'''
-            return '''«switchBlocks»    «expr»'''
-        }
+        if (isSwitchCond) return expr
         var blocks = (ifCondBlocks.isEmpty()) ? "" : '''    «FOR arg : ifCondBlocks»«arg»«ENDFOR»'''
         return '''«blocks»    return «expr»
         '''
@@ -286,74 +285,68 @@ class PythonExpressionGenerator {
     }
 
     private def String generateSwitchOperation(SwitchOperation expr, int ifLevel, boolean isLambda) {
-        // translate switch into a series of if / elif statements
         val attr = generateExpression(expr.argument, 0, isLambda)
-        val arg = expr.argument as RosettaSymbolReference
 
-        var funcNames = new ArrayList<String>()
+        var _thenFuncsBuilder = new StringConcatenation()
+        var _switchLogicBuilder= new StringConcatenation()
 
-        for (thenExpr : expr.cases) {
-            val thenExprDef = generateExpression(thenExpr.getExpression(), ifLevel + 1, isLambda)
-            val funcName = '''_then_«funcNames.size()+1»'''
-            funcNames.add(funcName)
-            switchCondBlocks.add(
-                '''
-                    def «funcName»():
-                        return «thenExprDef»
-                '''
-            )
-        }
-        
-        //default case
-        val defaultExprDef = generateExpression(expr.getDefault(), 0, isLambda)
-        val defaultFuncName = '''_then_default'''
-        funcNames.add(defaultFuncName)
-        switchCondBlocks.add(
-            '''
-                def «defaultFuncName»():
-                    return «defaultExprDef»
-            '''
-        )        
-        var _builder = new StringConcatenation()
-       
-        // Generate switch logic
         val indent = "    "
-        _builder.append("switchAttribute = ")
-        _builder.append(attr)
-        _builder.newLine()
-        // Append each conditional
-        for (i : 0 ..< expr.cases.size) {
-            val guard = expr.cases.get(i).getGuard()
-    
-            val prefix = (i == 0) ? "if " : "elif "
-            _builder.append(indent)
-            _builder.append(prefix)
-            if (guard.getLiteralGuard() !== null) {
-                val guardExpr = generateExpression(guard.getLiteralGuard(), 0, isLambda)
-                _builder.append("switchAttribute == ")
-                _builder.append(guardExpr)
-            } else {
-                val guardExpr = getGuardExpression(guard, isLambda)
-                _builder.append(guardExpr)
+        isSwitchCond=true
+        
+        for (pair : expr.cases.indexed) {
+            val currentCase = pair.value as SwitchCaseOrDefault
+            val funcName= (currentCase.isDefault()) ? "_then_default" : "_then_"+ (pair.key+1)
+            val thenExprDef= (currentCase.isDefault()) ? generateExpression(expr.getDefault(), 0, isLambda) : generateExpression(currentCase.getExpression(), ifLevel + 1, isLambda)
+            
+            _thenFuncsBuilder.append(indent)
+            _thenFuncsBuilder.append("def "+funcName + "():")
+            _thenFuncsBuilder.newLine
+            _thenFuncsBuilder.append(indent)
+            _thenFuncsBuilder.append("    return "+ thenExprDef)
+            _thenFuncsBuilder.newLine
+            
+            if(currentCase.isDefault()){
+                 // Default else
+                _switchLogicBuilder.append(indent)
+                _switchLogicBuilder.append("else:")
+                _switchLogicBuilder.newLine()
+                _switchLogicBuilder.append(indent)
+                _switchLogicBuilder.append("    return ")
+                _switchLogicBuilder.append(funcName)
+                _switchLogicBuilder.append("()")
+            }
+            else{
+                val guard =currentCase.getGuard()
+        
+                val prefix = (pair.key == 0) ? "if " : "elif "
+                _switchLogicBuilder.append(indent)
+                _switchLogicBuilder.append(prefix)
+                if (guard.getLiteralGuard() !== null) {
+                    val guardExpr = generateExpression(guard.getLiteralGuard(), 0, isLambda)
+                    _switchLogicBuilder.append("switchAttribute == ")
+                    _switchLogicBuilder.append(guardExpr)
+                } else {
+                    val guardExpr = getGuardExpression(guard, isLambda)
+                    _switchLogicBuilder.append(guardExpr)
+                }
+                _switchLogicBuilder.append(":")
+                _switchLogicBuilder.newLine()
+                _switchLogicBuilder.append(indent)
+                _switchLogicBuilder.append("    return ")
+                _switchLogicBuilder.append(funcName)
+                _switchLogicBuilder.append("()")
+                _switchLogicBuilder.newLine()
             }
-            _builder.append(":")
-            _builder.newLine()
-            _builder.append(indent)
-            _builder.append("    return ")
-            _builder.append(funcNames.get(i))
-            _builder.append("()")
-            _builder.newLine()
         }
-    
-        // Default else
-        _builder.append(indent)
-        _builder.append("else:")
-        _builder.newLine()
+
+        val _builder= new StringConcatenation
+        _builder.append(_thenFuncsBuilder.toString)
         _builder.append(indent)
-        _builder.append("    return ")
-        _builder.append(funcNames.last)
-        _builder.append("()")
-    
+        _builder.append("switchAttribute = ")
+        _builder.append(attr)
+        _builder.newLine
+        _builder.append(_switchLogicBuilder.toString)
+
         return _builder.toString
     }