📖 Command-line options support using docopt.

alexjc · alexjc · commit a4b2a65df946 · 2019-09-17T00:17:19.000+02:00
diff --git a/pytrain/__main__.py b/pytrain/__main__.py
@@ -1,13 +1,25 @@
 # PyTrain — Copyright (c) 2019, Alex J. Champandard.
+"""
+Usage: pytrain [-k FILTER] [-r ROOTDIR]
+
+Options:
+  -k FILTER         Select which tests to run by matching this substring filter.
+  -r ROOTDIR        Base directory from which to collect all the tasks.
+"""
 
 import asyncio
 
+from docopt import docopt
+
+from . import __version__
 from .application import Application
 from .registry import Registry
 
 
 def main():
-    registry = Registry()
+    config = docopt(__doc__, version=f"pytrain {__version__}")
+
+    registry = Registry(config)
     registry.load()
 
     loop = asyncio.get_event_loop()
diff --git a/pytrain/registry.py b/pytrain/registry.py
@@ -16,17 +16,19 @@ def config(self, key, default):
 
 
 class Registry:
-    def __init__(self):
+    def __init__(self, config):
         self.functions = []
         self.modules = []
 
+        self.config = config
+
     def load(self):
-        for root, _, files in os.walk("examples"):
+        for root, _, files in os.walk(self.config.get("-r") or "train"):
             if root.endswith("__pycache__"):
                 continue
 
             for filename in files:
-                if not filename.startswith("train_"):
+                if not filename.startswith(self.config.get("-k") or "train_"):
                     continue
 
                 basename = os.path.splitext(filename)[0]
