From ff428fd6f78de223ceea377bec663a75d9a0ec96 Mon Sep 17 00:00:00 2001
From: Dmitry Trofimov <dmitry.trofimov@jetbrains.com>
Date: Fri, 14 Dec 2012 18:26:32 +0100
Subject: [PATCH] Add path completion to os.path.join elements.

---
 .../python/psi/impl/PyPathEvaluator.java      | 50 ++++++++++---------
 1 file changed, 27 insertions(+), 23 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyPathEvaluator.java b/python/src/com/jetbrains/python/psi/impl/PyPathEvaluator.java
index 2adb08df2c81..3b0acc5cd7f2 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyPathEvaluator.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyPathEvaluator.java
@@ -28,7 +28,7 @@ public class PyPathEvaluator {
   }
 
   @Nullable
-  public static String evaluate(PyExpression expr, String path) {
+  public static String evaluate(PyExpression expr, String containingFilePath) {
     if (expr == null) {
       return null;
     }
@@ -36,27 +36,15 @@ public class PyPathEvaluator {
       final PyCallExpression call = (PyCallExpression)expr;
       final PyExpression[] args = call.getArguments();
       if (call.isCalleeText(PyNames.DIRNAME) && args.length == 1) {
-        String argValue = evaluate(args[0], path);
+        String argValue = evaluate(args[0], containingFilePath);
         return argValue == null ? null : new File(argValue).getParent();
       }
       else if (call.isCalleeText(PyNames.JOIN) && args.length >= 1) {
-        String result = null;
-        for (int i = 0; i<args.length; i++) {
-          String arg = evaluate(args[i], path);
-          if (arg == null) {
-            return null;
-          }
-          if (result == null) {
-            result = arg;
-          } else {
-            result = new File(result, arg).getPath();
-          }
-        }
-        return result;
+        return evaluatePathInJoin(containingFilePath, args, args.length);
       }
       else if (call.isCalleeText(PyNames.ABSPATH) && args.length == 1) {
-        String argValue = evaluate(args[0], path);
-        // relative to directory of 'path', not file
+        String argValue = evaluate(args[0], containingFilePath);
+        // relative to directory of 'containingFilePath', not file
         if (argValue == null) {
           return null;
         }
@@ -64,17 +52,17 @@ public class PyPathEvaluator {
           return argValue;
         }
         else {
-          return new File(new File(path).getParent(), argValue).getPath();
+          return new File(new File(containingFilePath).getParent(), argValue).getPath();
         }
       }
       else if (call.isCalleeText(PyNames.REPLACE) && args.length == 2) {
         final PyExpression callee = call.getCallee();
         if (!(callee instanceof PyQualifiedExpression)) return null;
         final PyExpression qualifier = ((PyQualifiedExpression)callee).getQualifier();
-        String result = evaluate(qualifier, path);
+        String result = evaluate(qualifier, containingFilePath);
         if (result == null) return null;
-        String arg1 = evaluate(args[0], path);
-        String arg2 = evaluate(args[1], path);
+        String arg1 = evaluate(args[0], containingFilePath);
+        String arg2 = evaluate(args[1], containingFilePath);
         if (arg1 == null || arg2 == null) return null;
         return result.replace(arg1, arg2);
       }
@@ -83,14 +71,14 @@ public class PyPathEvaluator {
       if (((PyReferenceExpression)expr).getQualifier() == null) {
         final String refName = ((PyReferenceExpression)expr).getReferencedName();
         if (PyNames.FILE.equals(refName)) {
-          return path;
+          return containingFilePath;
         }
         PsiElement result = ((PyReferenceExpression)expr).getReference(PyResolveContext.noImplicits()).resolve();
         if (result instanceof PyTargetExpression) {
           result = ((PyTargetExpression)result).findAssignedValue();
         }
         if (result instanceof PyExpression) {
-          return evaluate((PyExpression)result, path);
+          return evaluate((PyExpression)result, containingFilePath);
         }
       }
     }
@@ -99,4 +87,20 @@ public class PyPathEvaluator {
     }
     return null;
   }
+
+  public static String evaluatePathInJoin(String containingFilePath, PyExpression[] args, int endElement) {
+    String result = null;
+    for (int i = 0; i< endElement; i++) {
+      String arg = evaluate(args[i], containingFilePath);
+      if (arg == null) {
+        return null;
+      }
+      if (result == null) {
+        result = arg;
+      } else {
+        result = new File(result, arg).getPath();
+      }
+    }
+    return result;
+  }
 }