attempt to modularise

LICENSE.txt

@@ -1,4 +1,5 @@
 Copyright (c) 2015 Min RK, Florian Rathgeber, Michael McNeil Forbes
+2019 Casper da Costa-Luis
 
 Permission is hereby granted, free of charge, to any person obtaining
 a copy of this software and associated documentation files (the

nbstripout/__init__.py

@@ -0,0 +1,4 @@
+from ._nbstripout import install, uninstall, status, main
+from ._utils import pop_recursive, strip_output
+__all__ = ["install", "uninstall", "status", "main",
+ "pop_recursive", "strip_output"]

nbstripout.py → nbstripout/_nbstripout.py

@@ -85,11 +85,12 @@
 
  *.ipynb diff=ipynb
 """
-
 from __future__ import print_function
 from argparse import ArgumentParser, RawDescriptionHelpFormatter
 import io
 import sys
+from nbstripout._utils import strip_output
+__all__ = ["install", "uninstall", "status", "main"]
 
 input_stream = None
 if sys.version_info < (3, 0):
@@ -132,106 +133,6 @@ def write(nb, f):
  return current.write(nb, f, 'json')
 
 
-def _cells(nb):
- """Yield all cells in an nbformat-insensitive manner"""
- if nb.nbformat < 4:
- for ws in nb.worksheets:
- for cell in ws.cells:
- yield cell
- else:
- for cell in nb.cells:
- yield cell
-
-
-def pop_recursive(d, key, default=None):
- """dict.pop(key) where `key` is a `.`-delimited list of nested keys.
-
- >>> d = {'a': {'b': 1, 'c': 2}}
- >>> pop_recursive(d, 'a.c')
- 2
- >>> d
- {'a': {'b': 1}}
- """
- nested = key.split('.')
- current = d
- for k in nested[:-1]:
- if hasattr(current, 'get'):
- current = current.get(k, {})
- else:
- return default
- if not hasattr(current, 'pop'):
- return default
- return current.pop(nested[-1], default)
-
-
-def strip_output(nb, keep_output, keep_count, extra_keys=''):
- """
- Strip the outputs, execution count/prompt number and miscellaneous
- metadata from a notebook object, unless specified to keep either the outputs
- or counts.
-
- `extra_keys` could be 'metadata.foo cell.metadata.bar metadata.baz'
- """
- extra_keys = extra_keys.split()
- keys = {'metadata': [], 'cell': {'metadata': []}}
- for key in extra_keys:
- if key.startswith('metadata.'):
- keys['metadata'].append(key[len('metadata.'):])
- elif key.startswith('cell.metadata.'):
- keys['cell']['metadata'].append(key[len('cell.metadata.'):])
- else:
- sys.stderr.write('ignoring extra key `%s`' % key)
-
- nb.metadata.pop('signature', None)
- nb.metadata.pop('widgets', None)
- for field in keys['metadata']:
- pop_recursive(nb.metadata, field)
-
- for cell in _cells(nb):
-
- keep_output_this_cell = keep_output
-
- # Keep the output for these cells, but strip count and metadata
- if cell.metadata.get('init_cell') or cell.metadata.get('keep_output'):
- keep_output_this_cell = True
-
- # Remove the outputs, unless directed otherwise
- if 'outputs' in cell:
-
- # Default behavior strips outputs. With all outputs stripped,
- # there are no counts to keep and keep_count is ignored.
- if not keep_output_this_cell:
- cell['outputs'] = []
-
- # If keep_output_this_cell, but not keep_count, strip the counts
- # from the output.
- if keep_output_this_cell and not keep_count:
- for output in cell['outputs']:
- if 'execution_count' in output:
- output['execution_count'] = None
-
- # If keep_output_this_cell and keep_count, do nothing.
-
- # Remove the prompt_number/execution_count, unless directed otherwise
- if 'prompt_number' in cell and not keep_count:
- cell['prompt_number'] = None
- if 'execution_count' in cell and not keep_count:
- cell['execution_count'] = None
-
- # Always remove this metadata
- for output_style in ['collapsed', 'scrolled']:
- if output_style in cell.metadata:
- cell.metadata[output_style] = False
- if 'metadata' in cell:
- for field in ['collapsed', 'scrolled', 'ExecuteTime']:
- cell.metadata.pop(field, None)
- for (extra, fields) in keys['cell'].items():
- if extra in cell:
- for field in fields:
- pop_recursive(getattr(cell, extra), field)
- return nb
-
-
 def install(attrfile=None):
  """Install the git filter and set the git attributes."""
  from os import name, path

nbstripout/_utils.py

@@ -0,0 +1,102 @@
+import sys
+
+__all__ = ["pop_recursive", "strip_output"]
+
+
+def pop_recursive(d, key, default=None):
+ """dict.pop(key) where `key` is a `.`-delimited list of nested keys.
+
+ >>> d = {'a': {'b': 1, 'c': 2}}
+ >>> pop_recursive(d, 'a.c')
+ 2
+ >>> d
+ {'a': {'b': 1}}
+ """
+ nested = key.split('.')
+ current = d
+ for k in nested[:-1]:
+ if hasattr(current, 'get'):
+ current = current.get(k, {})
+ else:
+ return default
+ if not hasattr(current, 'pop'):
+ return default
+ return current.pop(nested[-1], default)
+
+
+def _cells(nb):
+ """Yield all cells in an nbformat-insensitive manner"""
+ if nb.nbformat < 4:
+ for ws in nb.worksheets:
+ for cell in ws.cells:
+ yield cell
+ else:
+ for cell in nb.cells:
+ yield cell
+
+
+def strip_output(nb, keep_output, keep_count, extra_keys=''):
+ """
+ Strip the outputs, execution count/prompt number and miscellaneous
+ metadata from a notebook object, unless specified to keep either the outputs
+ or counts.
+
+ `extra_keys` could be 'metadata.foo cell.metadata.bar metadata.baz'
+ """
+ extra_keys = extra_keys.split()
+ keys = {'metadata': [], 'cell': {'metadata': []}}
+ for key in extra_keys:
+ if key.startswith('metadata.'):
+ keys['metadata'].append(key[len('metadata.'):])
+ elif key.startswith('cell.metadata.'):
+ keys['cell']['metadata'].append(key[len('cell.metadata.'):])
+ else:
+ sys.stderr.write('ignoring extra key `%s`' % key)
+
+ nb.metadata.pop('signature', None)
+ nb.metadata.pop('widgets', None)
+ for field in keys['metadata']:
+ pop_recursive(nb.metadata, field)
+
+ for cell in _cells(nb):
+ keep_output_this_cell = keep_output
+
+ # Keep the output for these cells, but strip count and metadata
+ if cell.metadata.get('init_cell') or cell.metadata.get('keep_output'):
+ keep_output_this_cell = True
+
+ # Remove the outputs, unless directed otherwise
+ if 'outputs' in cell:
+
+ # Default behavior strips outputs. With all outputs stripped,
+ # there are no counts to keep and keep_count is ignored.
+ if not keep_output_this_cell:
+ cell['outputs'] = []
+
+ # If keep_output_this_cell, but not keep_count, strip the counts
+ # from the output.
+ if keep_output_this_cell and not keep_count:
+ for output in cell['outputs']:
+ if 'execution_count' in output:
+ output['execution_count'] = None
+
+ # If keep_output_this_cell and keep_count, do nothing.
+
+ # Remove the prompt_number/execution_count, unless directed otherwise
+ if 'prompt_number' in cell and not keep_count:
+ cell['prompt_number'] = None
+ if 'execution_count' in cell and not keep_count:
+ cell['execution_count'] = None
+
+ # Always remove this metadata
+ for output_style in ['collapsed', 'scrolled']:
+ if output_style in cell.metadata:
+ cell.metadata[output_style] = False
+ if 'metadata' in cell:
+ for field in ['collapsed', 'scrolled', 'ExecuteTime']:
+ cell.metadata.pop(field, None)
+ for (extra, fields) in keys['cell'].items():
+ if extra in cell:
+ for field in fields:
+ pop_recursive(getattr(cell, extra), field)
+ return nb

setup.py

@@ -1,4 +1,4 @@
-from setuptools import setup
+from setuptools import setup, find_packages
 
 with open('README.rst') as f:
  long_description = f.read()
@@ -32,10 +32,11 @@
 
  description='Strips outputs from Jupyter and IPython notebooks',
  long_description=long_description,
- py_modules=['nbstripout'],
+ packages=find_packages(),
+ provides=['nbstripout'],
  entry_points={
  'console_scripts': [
- 'nbstripout = nbstripout:main'
+ 'nbstripout = nbstripout._nbstripout:main'
  ],
  },
 

tests/test_utils.py

@@ -0,0 +1,7 @@
+from nbstripout._utils import pop_recursive
+
+
+def test_pop_recursive():
+ d = {'a': {'b': 1, 'c': 2}}
+ assert pop_recursive(d, 'a.c') == 2
+ assert d == {'a': {'b': 1}}