fix(parse_file): handle EnumDeclaration and AnnotationDeclaration in get_classes_and_methods

src/migration_bench/lang/java/eval/parse_file.py

@@ -7,7 +7,9 @@
 from migration_bench.common import utils
 
 
+ANNOTATION = "annotation"
 CLASS = "class"
+ENUM = "enum"
 INTERFACE = "interface"
 METHOD = "method"
 
@@ -34,6 +36,12 @@ def get_classes_and_methods(filename: str, content: str = None, add_line: bool =
         elif isinstance(node, javalang.tree.InterfaceDeclaration):
             current_class = node.name
             nodes.append((INTERFACE, current_class))
+        elif isinstance(node, javalang.tree.EnumDeclaration):
+            current_class = node.name
+            nodes.append((ENUM, current_class))
+        elif isinstance(node, javalang.tree.AnnotationDeclaration):
+            current_class = node.name
+            nodes.append((ANNOTATION, current_class))
         elif isinstance(node, javalang.tree.MethodDeclaration):
             if not current_class:
                 raise ValueError(

src/migration_bench/lang/java/eval/test_parse_file.py

@@ -19,6 +19,8 @@
 
 FILE_01 = "../native/src/main/java/qct/XmlBeautifier.java"
 FILE_02 = "./testdata/LabelManager.java"
+FILE_03 = "./testdata/Status.java"
+FILE_04 = "./testdata/Marker.java"
 
 
 class TestParseFile(unittest.TestCase):
@@ -78,6 +80,26 @@ class TestParseFile(unittest.TestCase):
                     ("method", ("get", False)),
                 ),
             ),
+            # Top-level enum with constants overriding an abstract method.
+            # Pre-fix: raised "Unable to find out class for method `describe`".
+            (
+                FILE_03,
+                {},
+                (
+                    ("enum", "Status"),
+                    ("method", ("describe", False)),
+                    ("method", ("describe", False)),
+                    ("method", ("describe", False)),
+                ),
+            ),
+            # Top-level annotation declaration.
+            (
+                FILE_04,
+                {},
+                (
+                    ("annotation", "Marker"),
+                ),
+            ),
         )
     )
     def test_get_classes_and_methods(self, filename, kwargs, expected_nodes):

src/migration_bench/lang/java/eval/testdata/Marker.java

@@ -0,0 +1,10 @@
+package com.example;
+
+/**
+ * Top-level annotation declaration with a method element. Previously crashed
+ * parse_file.get_classes_and_methods because AnnotationDeclaration was not
+ * handled.
+ */
+public @interface Marker {
+    String value() default "";
+}

src/migration_bench/lang/java/eval/testdata/Status.java

@@ -0,0 +1,23 @@
+package com.example;
+
+/**
+ * Top-level enum with constants whose bodies override an abstract method.
+ * This shape (seen in repos like thomasmueller/tinyStats) previously crashed
+ * parse_file.get_classes_and_methods because EnumDeclaration was not handled.
+ */
+public enum Status {
+    OK {
+        @Override
+        public String describe() {
+            return "ok";
+        }
+    },
+    BAD {
+        @Override
+        public String describe() {
+            return "bad";
+        }
+    };
+
+    public abstract String describe();
+}