refactor(mode): Use Enum for Mode

hardcoretech · May 16, 2022 · 515a622 · 515a622
1 parent 6bed7bd
commit 515a622
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 11 deletions.
diff --git a/src/missing.py b/src/missing.py
@@ -1,5 +1,6 @@
 #! /usr/bin/env python
 import argparse
+import enum
 import logging
 import os
 import re
@@ -9,6 +10,13 @@
 
 logger = logging.getLogger('missing')
 
+class Mode(enum.Enum):
+    ALL = 'all'
+    OBEY_GITIGNORE = 'obey_gitignore'
+    STAGED_ONLY = 'staged_only'
+
+    def __str__(self):
+        return self.value
 
 class Missing:
     RE_TEST_PY = re.compile(r'^test.*?\.py$')
@@ -96,9 +104,9 @@ def _is_tracked(self, path):
         return path in self._tracked_files
 
     def _find_tracked_files(self, mode):
-        if mode == 'obey_gitignore':
+        if Mode(mode) == Mode.OBEY_GITIGNORE:
             return self._list_files_obey_gitignore()
-        elif mode == 'staged_only':
+        elif Mode(mode) == Mode.STAGED_ONLY:
             return self._list_files_staged_only()
         else:
             return None
@@ -157,13 +165,13 @@ def main() -> int:
         '-q', '--quite', action='store_true', default=False, help='disable all log'
     )
 
+
     parser.add_argument(
         '-m',
         '--mode',
-        type=str,
-        default='all',
-        choices=['all', 'obey_gitignore', 'staged_only'],
-        help='Search scope',
+        type=Mode,
+        default=Mode.ALL,
+        choices=list(Mode)
     )
     args = parser.parse_args()
 

diff --git a/tests/test_missing.py b/tests/test_missing.py
@@ -3,7 +3,7 @@
 
 import pytest
 
-from missing import Missing
+from missing import Missing, Mode
 
 
 @pytest.fixture
@@ -45,7 +45,7 @@ def test__default_mode(self, temp_git_folder):
         assert os.path.isfile('tests/untracked/__init__.py')
         assert os.path.isfile('tests/staged/__init__.py')
         assert os.path.isfile('tests/__init__.py')
-        assert 0 == Missing(mode='all', exclude=[]).run()
+        assert 0 == Missing(mode=Mode.ALL, exclude=[]).run()
 
     def test__obey_gitignore(self, temp_git_folder):
         assert 1 == Missing(mode='obey_gitignore', exclude=[]).run()
@@ -54,7 +54,7 @@ def test__obey_gitignore(self, temp_git_folder):
         assert os.path.isfile('tests/untracked/__init__.py')
         assert os.path.isfile('tests/staged/__init__.py')
         assert os.path.isfile('tests/__init__.py')
-        assert 0 == Missing(mode='obey_gitignore', exclude=[]).run()
+        assert 0 == Missing(mode=Mode.OBEY_GITIGNORE, exclude=[]).run()
 
     def test__staged_only(self, temp_git_folder):
         assert 1 == Missing(mode='staged_only', exclude=[]).run()
@@ -63,7 +63,7 @@ def test__staged_only(self, temp_git_folder):
         assert not os.path.exists('tests/untracked/__init__.py')
         assert os.path.isfile('tests/staged/__init__.py')
         assert os.path.isfile('tests/__init__.py')
-        assert 0 == Missing(mode='staged_only', exclude=[]).run()
+        assert 0 == Missing(mode=Mode.STAGED_ONLY, exclude=[]).run()
 
     def test__exclude(self, temp_git_folder):
-        assert 0 == Missing(mode='all', exclude=['tests']).run()
+        assert 0 == Missing(mode=Mode.ALL, exclude=['tests']).run()