From 9d05f0aec2b024f83063e5209b2830c0c1a93d03 Mon Sep 17 00:00:00 2001
From: mattirn <matti.rintanikkola@gmail.com>
Date: Mon, 5 Apr 2021 15:21:34 +0200
Subject: [PATCH] Groovy REPL: add support for Groovy traits

---
 .../java/org/jline/script/GroovyEngine.java   | 38 +++++++++++++------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/groovy/src/main/java/org/jline/script/GroovyEngine.java b/groovy/src/main/java/org/jline/script/GroovyEngine.java
index 87d8d0913..c3441a992 100644
--- a/groovy/src/main/java/org/jline/script/GroovyEngine.java
+++ b/groovy/src/main/java/org/jline/script/GroovyEngine.java
@@ -76,6 +76,8 @@ public enum Format {JSON, GROOVY, NONE}
                                            , Pattern.DOTALL);
     private static final Pattern PATTERN_CLASS_DEF = Pattern.compile("^class\\s+(" + REGEX_VAR + ") .*?\\{.*?}(|\n)$"
                                                                   , Pattern.DOTALL);
+    private static final Pattern PATTERN_TRAIT_DEF = Pattern.compile("^trait\\s+" + REGEX_VAR + "\\s*\\{.*?}(|\n)$"
+                                                                    , Pattern.DOTALL);
     private static final Pattern PATTERN_CLASS = Pattern.compile("(.*?)\\.([A-Z_].*)");
     private static final String REGEX_PACKAGE = "([a-z][a-z_0-9]*\\.)*";
     private static final String REGEX_CLASS_NAME = "[A-Z_](\\w)*";
@@ -304,8 +306,11 @@ private static Object executeStatement(GroovyShell shell, Map<String, String> im
                 if (file.exists()) {
                     try {
                         shell.evaluate(file);
-                    } catch (MissingMethodExceptionNoStack ignore) {
-
+                    } catch (GroovyRuntimeException e) {
+                        if (!(e instanceof MissingMethodExceptionNoStack)              // thrown when class without main()
+                                && !(e.getCause() instanceof NoSuchMethodException)) { // thrown traits... no constructor
+                            throw e;
+                        }
                     }
                     if (importStatement) {
                         statement = importClass + matcher.group(4) + ".class";
@@ -322,7 +327,7 @@ private static Object executeStatement(GroovyShell shell, Map<String, String> im
             e.append(entry.getValue()).append("\n");
         }
         e.append(statement);
-        if (classDef(statement)) {
+        if (classOrTraitDef(statement)) {
             e.append("; null");
         }
         return shell.evaluate(e.toString());
@@ -375,8 +380,8 @@ private boolean functionDef(String statement) throws Exception{
         return out;
     }
 
-    private static boolean classDef(String statement) {
-        return PATTERN_CLASS_DEF.matcher(statement).matches();
+    private static boolean classOrTraitDef(String statement) {
+        return PATTERN_CLASS_DEF.matcher(statement).matches() || PATTERN_TRAIT_DEF.matcher(statement).matches();
     }
 
     private void refreshNameClass() {
@@ -471,7 +476,7 @@ protected static <T>T groovyOption(Map<String,Object> options, String option, T
 
     private Completer compileCompleter() {
         List<Completer> completers = new ArrayList<>();
-        completers.add(new ArgumentCompleter(new StringsCompleter("class", "print", "println"), NullCompleter.INSTANCE));
+        completers.add(new ArgumentCompleter(new StringsCompleter("class", "print", "println", "trait"), NullCompleter.INSTANCE));
         completers.add(new ArgumentCompleter(new StringsCompleter("def"), new StringsCompleter(methods::keySet)
                                            , NullCompleter.INSTANCE));
         completers.add(new ArgumentCompleter(new StringsCompleter("import")
@@ -1182,11 +1187,14 @@ public Inspector(GroovyEngine groovyEngine) {
                     sharedData.setVariable(entry.getKey(), execute("{" + m.group(1) + "->" + m.group(2) + "}"));
                 }
             }
-            for (String c : imports.keySet()) {
+            for (Map.Entry<String,String> entry : imports.entrySet()) {
                 try {
-                    addToNameClass(c, nameClass, shell);
-                } catch (Exception ignore) {
-
+                    executeStatement(shell, new HashMap<>(), entry.getValue());
+                    addToNameClass(entry.getKey(), nameClass, shell);
+                } catch (Exception e) {
+                    if (Log.isDebugEnabled()) {
+                        e.printStackTrace();
+                    }
                 }
             }
         }
@@ -1201,7 +1209,10 @@ public Class<?> evaluateClass(String objectStatement) {
                 involvedObject = execute(objectStatement);
                 out = involvedObject.getClass();
             } catch (Exception e) {
-                // ignore
+                Log.debug("objectStatement: ", objectStatement);
+                if (Log.isDebugEnabled()) {
+                    e.printStackTrace();
+                }
             }
             try {
                 if (out == null || out == Class.class) {
@@ -1231,7 +1242,10 @@ public Object execute(String statement) {
             Object out = null;
             try {
                 out = executeStatement(shell, imports, statement);
-            } catch (Exception ignore) {
+            } catch (Exception e) {
+                if (Log.isDebugEnabled()) {
+                    e.printStackTrace();
+                }
             } finally {
                 System.setOut(origOut);
                 System.setErr(origErr);