From 1179fc82563b6b5c5251b2c288cb094a4a94bc8f Mon Sep 17 00:00:00 2001
From: Jeff Zhang <zjffdu@apache.org>
Date: Fri, 12 Jan 2018 11:45:16 +0800
Subject: [PATCH] ZPEPELIN-3136. IPython Code Completion Improvement

---
 .../org/apache/zeppelin/python/IPythonInterpreter.java |  9 +++++++--
 .../apache/zeppelin/python/IPythonInterpreterTest.java | 10 +++++-----
 .../apache/zeppelin/spark/IPySparkInterpreterTest.java |  2 +-
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/python/src/main/java/org/apache/zeppelin/python/IPythonInterpreter.java b/python/src/main/java/org/apache/zeppelin/python/IPythonInterpreter.java
index 2184c1f10d5..5ae03f0230c 100644
--- a/python/src/main/java/org/apache/zeppelin/python/IPythonInterpreter.java
+++ b/python/src/main/java/org/apache/zeppelin/python/IPythonInterpreter.java
@@ -336,14 +336,19 @@ public int getProgress(InterpreterContext context) {
   @Override
   public List<InterpreterCompletion> completion(String buf, int cursor,
                                                 InterpreterContext interpreterContext) {
+    LOGGER.debug("Call completion for: " + buf);
     List<InterpreterCompletion> completions = new ArrayList<>();
     CompletionResponse response =
         ipythonClient.complete(
             CompletionRequest.getDefaultInstance().newBuilder().setCode(buf)
                 .setCursor(cursor).build());
     for (int i = 0; i < response.getMatchesCount(); i++) {
-      completions.add(new InterpreterCompletion(
-          response.getMatches(i), response.getMatches(i), ""));
+      String match = response.getMatches(i);
+      int lastIndexOfDot = match.lastIndexOf(".");
+      if (lastIndexOfDot != -1) {
+        match = match.substring(lastIndexOfDot + 1);
+      }
+      completions.add(new InterpreterCompletion(match, match, ""));
     }
     return completions;
   }
diff --git a/python/src/test/java/org/apache/zeppelin/python/IPythonInterpreterTest.java b/python/src/test/java/org/apache/zeppelin/python/IPythonInterpreterTest.java
index 6451aa3f418..d89ddac4469 100644
--- a/python/src/test/java/org/apache/zeppelin/python/IPythonInterpreterTest.java
+++ b/python/src/test/java/org/apache/zeppelin/python/IPythonInterpreterTest.java
@@ -196,9 +196,9 @@ public static void testInterpreter(final Interpreter interpreter) throws IOExcep
     context = getInterpreterContext();
     completions = interpreter.completion("sys.std", 7, context);
     assertEquals(3, completions.size());
-    assertEquals("sys.stderr", completions.get(0).getValue());
-    assertEquals("sys.stdin", completions.get(1).getValue());
-    assertEquals("sys.stdout", completions.get(2).getValue());
+    assertEquals("stderr", completions.get(0).getValue());
+    assertEquals("stdin", completions.get(1).getValue());
+    assertEquals("stdout", completions.get(2).getValue());
 
     // there's no completion for 'a.' because it is not recognized by compiler for now.
     context = getInterpreterContext();
@@ -227,14 +227,14 @@ public static void testInterpreter(final Interpreter interpreter) throws IOExcep
     st = "a.co";
     completions = interpreter.completion(st, st.length(), context);
     assertEquals(1, completions.size());
-    assertEquals("a.count", completions.get(0).getValue());
+    assertEquals("count", completions.get(0).getValue());
 
     // cursor is in the middle of code
     context = getInterpreterContext();
     st = "a.co\b='hello";
     completions = interpreter.completion(st, 4, context);
     assertEquals(1, completions.size());
-    assertEquals("a.count", completions.get(0).getValue());
+    assertEquals("count", completions.get(0).getValue());
 
     // ipython help
     context = getInterpreterContext();
diff --git a/spark/src/test/java/org/apache/zeppelin/spark/IPySparkInterpreterTest.java b/spark/src/test/java/org/apache/zeppelin/spark/IPySparkInterpreterTest.java
index 54922745365..d2b01ce7986 100644
--- a/spark/src/test/java/org/apache/zeppelin/spark/IPySparkInterpreterTest.java
+++ b/spark/src/test/java/org/apache/zeppelin/spark/IPySparkInterpreterTest.java
@@ -162,7 +162,7 @@ public void run() {
     // completions
     List<InterpreterCompletion> completions = iPySparkInterpreter.completion("sc.ran", 6, getInterpreterContext());
     assertEquals(1, completions.size());
-    assertEquals("sc.range", completions.get(0).getValue());
+    assertEquals("range", completions.get(0).getValue());
 
     // pyspark streaming
     context = getInterpreterContext();