Fix improper behavior if multiple arguments are issued

1 files changed, 27 insertions, 12 deletions

diff --git a/purge_path/purge_path.py b/purge_path/purge_path.py
index 322925b..dab14e4 100755
--- a/purge_path/purge_path.py
+++ b/purge_path/purge_path.py
@@ -1,20 +1,35 @@
 #!/usr/bin/env python
+from __future__ import print_function
 import os
 import sys
- 
+
+
+def purge(path, pattern):
+    if not isinstance(path, list):
+        path = path.split(':')
+
+    path_new = ''
+    for rec in path:
+        if rec == pattern:
+            # ignore pattern
+            continue
+        path_new += rec + ':'
+
+    return path_new[:-1]
+
 
 def main():
-    ARGS = sys.argv[1:]
-    PATH = os.environ['PATH']
-    PATH_NEW = []
- 
-    for arg in ARGS:
-        for path in PATH.split(':'):
-            if arg in path:
-                continue
-            PATH_NEW.append(path)   
- 
-    print("{0}".format(":".join(PATH_NEW)))
+    args = sys.argv[1:]
+    path_orig = os.environ['PATH']
+    path_new = path_orig
+
+    for arg in args:
+        path_new = purge(path_new, arg)
+
+    if path_new:
+        print(path_new)
+    else:
+        print(path_orig)
 
 
 if __name__ in '__main__':