Support config files

nbstripout/_nbstripout.py

@@ -109,7 +109,7 @@
     *.ipynb diff=ipynb
 """
 
-from argparse import ArgumentParser, RawDescriptionHelpFormatter
+from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter
 import collections
 import io
 import json
@@ -120,7 +120,7 @@
 import sys
 import warnings
 
-from nbstripout._utils import strip_output, strip_zeppelin_output
+from nbstripout._utils import strip_output, strip_zeppelin_output, merge_configuration_file
 try:
     # Jupyter >= 4
     from nbformat import read, write, NO_CONVERT
@@ -351,7 +351,7 @@ def status(git_config, install_location=INSTALL_LOCATION_LOCAL, verbose=False):
         return 1
 
 
-def main():
+def setup_commandline() -> Namespace:
     parser = ArgumentParser(epilog=__doc__, formatter_class=RawDescriptionHelpFormatter)
     task = parser.add_mutually_exclusive_group()
     task.add_argument('--dry-run', action='store_true',
@@ -377,7 +377,7 @@ def main():
                         help='Space separated list of extra keys to strip '
                         'from metadata, e.g. metadata.foo cell.metadata.bar')
     parser.add_argument('--drop-empty-cells', action='store_true',
-                        help='Remove cells where `source` is empty or contains only whitepace')
+                        help='Remove cells where `source` is empty or contains only whitespace')
     parser.add_argument('--drop-tagged-cells', default='',
                         help='Space separated list of cell-tags that remove an entire cell')
     parser.add_argument('--strip-init-cells', action='store_true',
@@ -402,7 +402,13 @@ def main():
                         help='Prints stripped files to STDOUT')
 
     parser.add_argument('files', nargs='*', help='Files to strip output from')
-    args = parser.parse_args()
+
+    return parser
+
+
+def main():
+    parser = setup_commandline()
+    args = merge_configuration_file(parser)
 
     git_config = ['git', 'config']
 

nbstripout/_utils.py

@@ -1,5 +1,9 @@
-from collections import defaultdict
+from argparse import ArgumentParser, Namespace, _StoreTrueAction, _StoreFalseAction
+import os
 import sys
+from collections import defaultdict
+from functools import partial
+from typing import Any, Dict, Optional
 
 __all__ = ["pop_recursive", "strip_output", "strip_zeppelin_output", "MetadataError"]
 
@@ -153,3 +157,106 @@ def strip_output(nb, keep_output, keep_count, extra_keys=[], drop_empty_cells=Fa
         for field in keys['cell']:
             pop_recursive(cell, field)
     return nb
+
+
+def process_pyproject_toml(toml_file_path: str, parser: ArgumentParser) -> Optional[Dict[str, Any]]:
+    """Extract config mapping from pyproject.toml file."""
+    try:
+        import tomllib  # python 3.11+
+    except ModuleNotFoundError:
+        import tomli as tomllib
+
+    with open(toml_file_path, 'rb') as f:
+        dict_config = tomllib.load(f).get('tool', {}).get('nbstripout', None)
+
+    if not dict_config:
+        # could be {} if 'tool' not found, or None if 'nbstripout' not found
+        return dict_config
+
+    # special processing of boolean options, make sure we don't have invalid types
+    for a in parser._actions:
+        if a.dest in dict_config and isinstance(a, (_StoreTrueAction, _StoreFalseAction)):
+            if not isinstance(dict_config[a.dest], bool):
+                raise ValueError(f'Argument {a.dest} in pyproject.toml must be a boolean, not {dict_config[a.dest]}')
+
+    return dict_config
+
+
+def process_setup_cfg(cfg_file_path, parser: ArgumentParser) -> Optional[Dict[str, Any]]:
+    """Extract config mapping from setup.cfg file."""
+    import configparser
+
+    reader = configparser.ConfigParser()
+    reader.read(cfg_file_path)
+    if not reader.has_section('nbstripout'):
+        return None
+
+    raw_config = reader['nbstripout']
+    dict_config = dict(raw_config)
+
+    # special processing of boolean options, to convert various configparser bool types to true/false
+    for a in parser._actions:
+        if a.dest in raw_config and isinstance(a, (_StoreTrueAction, _StoreFalseAction)):
+            dict_config[a.dest] = raw_config.getboolean(a.dest)
+
+    return dict_config
+
+
+def merge_configuration_file(parser: ArgumentParser, args_str=None) -> Namespace:
+    """Merge flags from config files into args."""
+    CONFIG_FILES = {
+        'pyproject.toml': partial(process_pyproject_toml, parser=parser),
+        'setup.cfg': partial(process_setup_cfg, parser=parser),
+    }
+
+    # parse args as-is to look for configuration files
+    args = parser.parse_args(args_str)
+
+    # Traverse the file tree common to all files given as argument looking for
+    # a configuration file
+    # TODO: make this more like Black:
+    #       By default Black looks for pyproject.toml starting from the common base directory of all files and
+    #       directories passed on the command line. If it’s not there, it looks in parent directories. It stops looking
+    #       when it finds the file, or a .git directory, or a .hg directory, or the root of the file system, whichever
+    #       comes first.
+    # if no files are given, start from cwd
+    config_path = os.path.commonpath([os.path.abspath(file) for file in args.files]) if args.files else os.path.abspath(os.getcwd())
+    config: Optional[Dict[str, Any]] = None
+    while True:
+        for config_file, processor in CONFIG_FILES.items():
+            config_file_path = os.path.join(config_path, config_file)
+            if os.path.isfile(config_file_path):
+                config = processor(config_file_path)
+                if config is not None:
+                    break
+        if config is not None:
+            break
+        config_path, tail = os.path.split(config_path)
+        if not tail:
+            break
+
+    # black starts with default arguments (from click), updates that with the config file,
+    # then applies command line arguments. this all happens in the click framework, before main() is called
+    # we can use parser.set_defaults
+    if config:
+        # check all arguments are valid
+        valid_args = vars(args).keys()
+        for name in config.keys():
+            if name not in valid_args:
+                raise ValueError(f'{name} in the config file is not a valid option')
+
+        # separate into default-overrides and special treatment
+        extra_keys: Optional[str] = None
+        if 'extra_keys' in config:
+            extra_keys = config['extra_keys']
+            del config['extra_keys']
+
+        # merge the configuration options as new defaults, and re-parse the arguments
+        parser.set_defaults(**config)
+        args = parser.parse_args(args_str)
+
+        # merge extra_keys using set union
+        if extra_keys:
+            args.extra_keys = ' '.join(sorted(set(extra_keys.split()) | set(args.extra_keys.split())))
+
+    return args

tests/conftest.py

@@ -0,0 +1 @@
+pytest_plugins = "pytester"