FEAT(index) add --sources CLI to limit sources to run,

args can be both Source.name & Source's positional index.

Author: ankostis

src/promnesia/__main__.py

@@ -19,7 +19,19 @@
 from .extract import extract_visits, make_filter
 
 
-def iter_all_visits() -> Iterator[Res[DbVisit]]:
+def _decide_indexers(
+        indexers: Iterable[Source], sources_subset: Iterable[str]
+    ) -> Iterable[Source]:
+    unknown = set(sources_subset) - set(range(len(indexers))) - {i.name for i in indexers}
+    if unknown:
+        # TODO: raise special exception on bad-sources and handle it politely in CLI.
+        raise ValueError(f"Unknown source(s): {', '.join(str(i) for i in unknown)}")
+    # Substitute positions with names, to facilitate debugging.
+    named_indexers = {i.name: i for i in indexers}
+    return {named_indexers[i] if isinstance(i, str) else indexers[i] for i in sources_subset}
+
+
+def iter_all_visits(sources_subset: Iterable[str]=()) -> Iterator[Res[DbVisit]]:
     cfg = config.get()
     output_dir = cfg.output_dir
     # not sure if belongs here??
@@ -29,8 +41,16 @@ def iter_all_visits() -> Iterator[Res[DbVisit]]:
 
     hook = cfg.hook
 
-    indexers = cfg.sources
-    for idx in indexers:
+    indexers = list(cfg.sources)
+
+    if sources_subset:
+        indexers = _decide_indexers(indexers, sources_subset)
+    logger.info("sources to extract: %s", ", ".join(i.name for i in indexers))
+
+    for i, idx in enumerate(indexers):
+        if sources_subset and idx.name not in sources_subset and i not in sources_subset:
+            continue
+
         if isinstance(idx, Exception):
             yield idx
             continue
@@ -46,11 +66,11 @@ def iter_all_visits() -> Iterator[Res[DbVisit]]:
                     yield e
 
 
-def _do_index(dry: bool=False) -> Iterable[Exception]:
+def _do_index(dry: bool=False, sources_subset: Iterable[str]=()) -> Iterable[Exception]:
     # also keep & return errors for further display
     errors: List[Exception] = []
     def it() -> Iterable[Res[DbVisit]]:
-        for v in iter_all_visits():
+        for v in iter_all_visits(sources_subset):
             if isinstance(v, Exception):
                 errors.append(v)
             yield v
@@ -68,10 +88,10 @@ def it() -> Iterable[Res[DbVisit]]:
     return errors
 
 
-def do_index(config_file: Path, dry: bool=False) -> None:
+def do_index(config_file: Path, dry: bool=False, sources_subset: Iterable[str]=()) -> None:
     config.load_from(config_file) # meh.. should be cleaner
     try:
-        errors = list(_do_index(dry=dry))
+        errors = list(_do_index(dry=dry, sources_subset=sources_subset))
     finally:
         config.reset()
     if len(errors) > 0:
@@ -104,7 +124,14 @@ def inner(*args, **kwargs):
     return res
 
 
-def do_demo(*, index_as: str, params: Sequence[str], port: Optional[str], config_file: Optional[Path], name: str='demo') -> None:
+def do_demo(*
+        index_as: str,
+        params: Sequence[str],
+        port: Optional[str],
+        config_file: Optional[Path],
+        name: str='demo',
+        sources_subset: Iterable[str]=(),
+    ) -> None:
     from pprint import pprint
     with TemporaryDirectory() as tdir:
         outdir = Path(tdir)
@@ -120,7 +147,7 @@ def do_demo(*, index_as: str, params: Sequence[str], port: Optional[str], config
             )
             config.instance = cfg
 
-        errors = list(_do_index())
+        errors = list(_do_index(sources_subset=sources_subset))
         if len(errors) > 0:
             logger.error('%d errors during indexing (see logs above for backtraces)', len(errors))
         for e in errors:
@@ -245,6 +272,14 @@ def cli_doctor_server(args: argparse.Namespace) -> None:
     logger.info('You should see the database path and version above!')
 
 
+def _parse_ordinal_or_name(s: str):
+    try:
+        s = int(s)
+    except ValueError:
+        pass
+    return s
+
+
 def main() -> None:
     # TODO longer, literate description?
 
@@ -256,6 +291,8 @@ def main() -> None:
     ep.add_argument('--dry', action='store_true', help="Dry run, won't touch the database, only print the results out")
     # TODO use some way to override or provide config only via cmdline?
     ep.add_argument('--intermediate', required=False, help="Used for development, you don't need it")
+    ep.add_argument('--sources', required=False, action="extend", nargs="+", type=_parse_ordinal_or_name,
+                    help="Subset of source(s) to run (name or 0-indexed position);  use `promnisia --dry` to view sources")
 
     sp = subp.add_parser('serve', help='Serve a link database', formatter_class=F) # type: ignore
     server.setup_parser(sp)
@@ -276,6 +313,8 @@ def main() -> None:
         default='guess',
         help='Promnesia source to index as (see https://github.com/karlicoss/promnesia/tree/master/src/promnesia/sources for the full list)',
     )
+    ap.add_argument('--sources', required=False, action="extend", nargs="+", type=_parse_ordinal_or_name,
+                    help="Subset of source(s) to run (name or 0-indexed position);  use `promnisia --dry` to view sources")
     ap.add_argument('params', nargs='*', help='Optional extra params for the indexer')
 
     isp = subp.add_parser('install-server', help='Install server as a systemd service (for autostart)', formatter_class=F)
@@ -315,13 +354,20 @@ def main() -> None:
 
     with get_tmpdir() as tdir: # TODO??
         if args.mode == 'index':
-             do_index(config_file=args.config, dry=args.dry)
+             do_index(config_file=args.config, dry=args.dry, sources_subset=args.sources)
         elif args.mode == 'serve':
             server.run(args)
         elif args.mode == 'demo':
             # TODO not sure if 'as' is that useful
             # something like Telegram/Takeout is too hard to setup to justify adhoc mode like this?
-            do_demo(index_as=getattr(args, 'as'), params=args.params, port=args.port, config_file=args.config, name=args.name)
+            do_demo(
+                index_as=getattr(args, 'as'),
+                params=args.params,
+                port=args.port,
+                config_file=args.config,
+                name=args.name,
+                sources_subset=args.sources,
+                )
         elif args.mode == 'install-server': # todo rename to 'autostart' or something?
             install_server.install(args)
         elif args.mode == 'config':