Improve c/h flip

1 files changed, 18 insertions, 18 deletions

diff --git a/dotvimrc b/dotvimrc
index 083c95e..e0121fe 100644
--- a/dotvimrc
+++ b/dotvimrc
@@ -212,24 +212,24 @@ def loadfile(filename):
 
 currentfile = vim.eval('expand("%")')
 
-if currentfile.endswith(".c"):
-    loadfile(currentfile[:-2] + ".h")
-elif currentfile.endswith(".h"):
-    basename = currentfile[:-2]
-    if os.path.exists(basename + ".cpp"):
-        loadfile(basename + ".cpp")
-    else:
-        loadfile(basename + ".c")
-elif currentfile.endswith(".cpp"):
-    basename = currentfile[:-4]
-    if os.path.exists(basename + ".h"):
-        loadfile(basename + ".h")
-    else:
-        loadfile(basename + ".hpp")
-elif currentfile.endswith(".hpp"):
-    basename = currentfile[:-4]
-    loadfile(basename + ".cpp")
-
+c_exts = ["c", "cpp", "cc"]
+h_exts = ["h", "hpp"]
+
+currentfile_split = currentfile.split(".")
+if len(currentfile_split) > 0:
+    ext = currentfile_split[-1]
+    basename = currentfile[:-(1+len(ext))]
+
+    if ext in c_exts:
+        for h_ext in h_exts:
+            h_filename = basename + "." + h_ext
+            if os.path.exists(h_filename):
+                loadfile(h_filename)
+    elif ext in h_exts:
+        for c_ext in c_exts:
+            c_filename = basename + "." + c_ext
+            if os.path.exists(c_filename):
+                loadfile(c_filename)
 endpython
 endfun