Merge pull request #773 from gnanaprakash-ravi/ZEIssue230

Update Annotations to generate code
zinggAI · Jan 26, 2024 · f9e9528 · f9e9528
2 parents d5159cc + 62e2933
commit f9e9528
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 22 deletions.
diff --git a/common/py/src/main/java/zingg/common/py/processors/PythonClassProcessor.java b/common/py/src/main/java/zingg/common/py/processors/PythonClassProcessor.java
@@ -47,14 +47,14 @@ public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment
                 // __init__ method
                 System.out.println("    def __init__(self" +
                         generateConstructorParameters(classElement) + "):");
-                if (element.getSimpleName().contentEquals("Pipe")) {
-                    generateClassInitializationCode(classElement, element);
-                }
-                for (VariableElement field : ElementFilter.fieldsIn(classElement.getEnclosedElements())) {
-                    if (!field.getSimpleName().contentEquals("serialVersionUID")) {
-                        generateFieldInitializationCode(field, element);
-                    }
-                }
+                generateClassInitializationCode(classElement, element);
+
+                // for (VariableElement field : ElementFilter.fieldsIn(classElement.getEnclosedElements())) {
+                //     if (!field.getSimpleName().contentEquals("serialVersionUID")) {
+                //         generateFieldInitializationCode(field, element);
+                //     }
+                // }
+
                 for (ExecutableElement methodElement : ElementFilter.methodsIn(classElement.getEnclosedElements())) {
                     if (methodElement.getAnnotation(PythonMethod.class) != null) {
                         methodNames.add(methodElement.getSimpleName().toString());
@@ -87,17 +87,19 @@ private void generateImportsAndDeclarations() {
     }
 
     private void generateClassInitializationCode(TypeElement classElement, Element element) {
-        System.out.println("        self." + element.getSimpleName().toString().toLowerCase() + " = getJVM().zingg.spark.client.pipe.SparkPipe()");
+        if (element.getSimpleName().contentEquals("Pipe")) {
+            System.out.println("        self." + element.getSimpleName().toString().toLowerCase() + " = getJVM().zingg.spark.client.pipe.SparkPipe()");
+        }
     }
 
-    private void generateFieldInitializationCode(VariableElement field, Element element) {
-        String fieldName = field.getSimpleName().toString();
-        String fieldAssignment = "self." + element.getSimpleName().toString().toLowerCase() + "." + fieldName + " = " + fieldName;
+    // private void generateFieldInitializationCode(VariableElement field, Element element) {
+    //     String fieldName = field.getSimpleName().toString();
+    //     String fieldAssignment = "self." + element.getSimpleName().toString().toLowerCase() + "." + fieldName + " = " + fieldName;
 
-        if (!fieldName.startsWith("FORMAT_")) {
-            System.out.println("        " + fieldAssignment);
-        }
-    }
+    //     if (!fieldName.startsWith("FORMAT_")) {
+    //         System.out.println("        " + fieldAssignment);
+    //     }
+    // }
 
     private String generateConstructorParameters(TypeElement classElement) {
         StringBuilder parameters = new StringBuilder();

diff --git a/common/py/src/main/java/zingg/common/py/processors/PythonMethodProcessor.java b/common/py/src/main/java/zingg/common/py/processors/PythonMethodProcessor.java
@@ -7,7 +7,7 @@
 import javax.lang.model.type.TypeMirror;
 import javax.lang.model.type.TypeKind;
 import java.util.Set;
-import java.util.logging.Logger;
+// import java.util.logging.Logger;
 
 import javax.lang.model.element.*;
 import zingg.common.py.annotations.*;
@@ -38,7 +38,8 @@ public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment
                     if (methodNames.contains(methodElement.getSimpleName().toString())) {
                         // LOG.info("Generating Python method for: " + methodElement.getSimpleName());
                         System.out.println("    def " + methodElement.getSimpleName() + "(self" +
-                                generateMethodSignature(methodElement) + "):\n        " + generateMethodReturn(methodElement));
+                                generateMethodSignature(methodElement) + "):");
+                        generateMethodReturn(methodElement);
                         generateFieldAssignment(methodElement);
                     }
                 }
@@ -64,14 +65,15 @@ private String generateMethodParameters(ExecutableElement methodElement) {
         return parameters.toString();
     }
 
-    private String generateMethodReturn(ExecutableElement methodElement) {
+    private void generateMethodReturn(ExecutableElement methodElement) {
         TypeMirror returnType = methodElement.getReturnType();
         if (returnType.getKind() == TypeKind.VOID) {
-            return "";
+            return;
         } else {
             String returnTypeString = resolveType(returnType);
-            String variableName = methodElement.getSimpleName().toString();
-            return "return " + variableName;
+            String methodName = methodElement.getSimpleName().toString();
+            String className = methodElement.getEnclosingElement().getSimpleName().toString();
+            System.out.println("        return self." + className.toLowerCase() + "." + methodName + "()");
         }
     }