From 539bf7e293d88f8b9b77b67e4bb8191f801a887b Mon Sep 17 00:00:00 2001
From: Franz Beaune <franz@appleseedhq.net>
Date: Sun, 11 Sep 2016 11:20:15 +0200
Subject: [PATCH] Make sortincludes.py more flexible

---
 scripts/sortincludes.py | 39 ++++++++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/scripts/sortincludes.py b/scripts/sortincludes.py
index aa4a0c4ca9..67665630a6 100644
--- a/scripts/sortincludes.py
+++ b/scripts/sortincludes.py
@@ -32,6 +32,21 @@
 import sys
 
 
+#--------------------------------------------------------------------------------------------------
+# Utility functions.
+#--------------------------------------------------------------------------------------------------
+
+def walk(directory, recursive):
+    if recursive:
+        for dirpath, dirnames, filenames in os.walk(directory):
+            for filename in filenames:
+                yield os.path.join(dirpath, filename)
+    else:
+        dirpath, dirnames, filenames = os.walk(directory).next()
+        for filename in filenames:
+            yield os.path.join(dirpath, filename)
+
+
 #--------------------------------------------------------------------------------------------------
 # Processing code.
 #--------------------------------------------------------------------------------------------------
@@ -61,13 +76,6 @@ def process_file(filepath):
 
     os.remove(filepath)
     os.rename(filepath + ".processed", filepath)
-    
-def process_recursively():
-    for dirpath, dirnames, filenames in os.walk("."):
-        for filename in filenames:
-            ext = os.path.splitext(filename)[1]
-            if ext == ".h" or ext == ".cpp":
-                process_file(os.path.join(dirpath, filename))
 
 
 #--------------------------------------------------------------------------------------------------
@@ -78,19 +86,16 @@ def main():
     parser = argparse.ArgumentParser(description="sort #include clauses in c++ source code.")
     parser.add_argument("-r", "--recursive", action='store_true', dest='recursive',
                         help="process all files in the specified directory and all its subdirectories")
-    parser.add_argument("filepath", nargs="?", help="file to process")
+    parser.add_argument("path", help="file or directory to process")
     args = parser.parse_args()
 
-    if args.filepath:
-        if args.recursive:
-            print("cannot simultaneously specify a file path and use the --recursive flag.")
-            sys.exit(1)
-        print(args.filepath)
+    if os.path.isfile(args.path):
+        process_file(args.path)
     else:
-        if not args.recursive:
-            print("either a file path or the --recursive flag must be specified.")
-            sys.exit(1)
-        process_recursively()
+        for filepath in walk(args.path, args.recursive):
+            ext = os.path.splitext(filepath)[1]
+            if ext == ".h" or ext == ".cpp":
+                process_file(filepath)
 
 if __name__ == '__main__':
     main()