ENH: Added uncommitted changes support

1 files changed, 14 insertions, 6 deletions

diff --git a/tools/linter.py b/tools/linter.py
index f0c2b7927..509e199d8 100644
--- a/tools/linter.py
+++ b/tools/linter.py
@@ -16,15 +16,20 @@ class DiffLinter:
         self.branch = branch
         self.repo = Repo('.')
 
-    def get_branch_diff(self):
+    def get_branch_diff(self, uncommitted):
         commit = self.repo.merge_base(BASE_BRANCH, self.branch)[0]
-        diff = self.repo.git.diff(commit, self.branch, '*.py')
+
+        if uncommitted:
+            diff = self.repo.git.diff(self.branch, '***.py')
+        else:
+            diff = self.repo.git.diff(commit, self.branch, '***.py')
         return diff
 
     def run_pycodestyle(self, diff):
         """
             Original Author: Josh Wilson (@person142)
-            Source: https://github.com/scipy/scipy/blob/master/tools/lint_diff.py
+            Source:
+              https://github.com/scipy/scipy/blob/master/tools/lint_diff.py
             Run pycodestyle on the given diff.
         """
         res = subprocess.run(
@@ -35,18 +40,21 @@ class DiffLinter:
         )
         return res.returncode, res.stdout
 
-    def run_lint(self):
-        diff = self.get_branch_diff()
+    def run_lint(self, uncommitted):
+        diff = self.get_branch_diff(uncommitted)
         retcode, errors = self.run_pycodestyle(diff)
 
         errors and print(errors)
 
         sys.exit(retcode)
 
+
 if __name__ == '__main__':
     parser = ArgumentParser()
     parser.add_argument("--branch", type=str, default='master',
                         help="The branch to diff against")
+    parser.add_argument("--uncommitted", action='store_true',
+                        help="Check only uncommitted changes")
     args = parser.parse_args()
 
-    DiffLinter(args.branch).run_lint()
+    DiffLinter(args.branch).run_lint(args.uncommitted)