Update documentation and add more tests

README.rst

@@ -50,6 +50,12 @@ named `_merged.ipynb`,
 
     nbmerge --recursive -i -p ".*intro.*" -o _merged.ipynb
 
+Finally, you can also instruct the script to demarcate the boundary
+between each original file with the `-b` / `-boundary [BOUNDARY]` flag.
+The `src_nb` value in the metadata for the first cell in each original
+notebook will then contain the path of the original notebook, relative to
+the cwd at the point of script execution.
+
 Lineage
 =======
 
@@ -82,3 +88,4 @@ Right now, only the basic (originally fperez's) functionality is
 implemented. However, I'm going to follow 
 `kynan's <https://github.com/kynan>`__ lead and slowly pull in functionality
 similar to his ``nbstripout`` package.
+

nbmerge/__init__.py

@@ -164,8 +164,7 @@ def parse_plan(args=None, base_dir=None):
     file_paths = args.files[:]
     for file_path in file_paths:
         if not os.path.exists(file_path):
-            print("Notebook `{}` does not exist".format(file_path))
-            sys.exit(1)
+            raise IOError("Notebook `{}` does not exist".format(file_path))
 
     if args.recursive:
         # If you specify any files, they are added first, in order.
@@ -180,8 +179,8 @@ def parse_plan(args=None, base_dir=None):
             'verbose': args.verbose}
 
 
-def main():
-    plan = parse_plan()
+def main(args=None):
+    plan = parse_plan(args or sys.argv)
 
     nb = merge_notebooks(plan['base_dir'],
                          plan['notebooks'],

tests/test_merge.py

@@ -2,6 +2,7 @@
 import sys
 import unittest
 
+from tempfile import mkstemp
 from nbformat import reads
 from nbmerge import merge_notebooks, main, parse_plan, annotate_source_path
 
@@ -19,14 +20,33 @@ def file_names_from(file_paths):
 
 
 class TestMerge(unittest.TestCase):
+    def setUp(self):
+        if not hasattr(sys.stdout, "getvalue"):
+            self.fail("need to run in buffered mode")
+
     def _validate_merged_three(self, merged):
         self.assertEqual(len(merged.cells), 6)
         self.assertEqual(merged.metadata['test_meta']['title'], "Page 1")
         self.assertEqual(merged.metadata['final_answer'], 42)
 
-    def test_merge(self):
+    def test_merge_defaults(self):
         self._validate_merged_three(merge_notebooks(FIXTURES_DIR, TARGET_NBS))
 
+    def test_merge_verbose(self):
+        nb = merge_notebooks(FIXTURES_DIR, TARGET_NBS, verbose=True)
+        self._validate_merged_three(nb)
+        lines = sys.stdout.getvalue().splitlines()
+        self.assertEqual(lines[0].strip(), "Merging notebooks...")
+        for target, line in zip(TARGET_NBS, lines[1:4]):
+            self.assertEqual(line.strip(), "Reading `{}`".format(target))
+
+    def test_merge_with_boundary_key(self):
+        nb = merge_notebooks(FIXTURES_DIR, TARGET_NBS, boundary_key='xxx')
+        self._validate_merged_three(nb)
+        self.assertEqual(nb.cells[0].metadata['xxx'], '1_Intro.ipynb')
+        self.assertEqual(nb.cells[2].metadata['xxx'], '2_Middle.ipynb')
+        self.assertEqual(nb.cells[4].metadata['xxx'], '3_Conclusion.ipynb')
+
     def test_parse_plan(self):
         header_nb = os.path.join(FIXTURES_DIR, "Header.ipynb")
         plan = parse_plan(["-o", "myfile.ipynb",
@@ -39,6 +59,8 @@ def test_parse_plan(self):
                           "1_Intro_In_Sub.ipynb", "2_Middle.ipynb"])
         self.assertTrue(plan["verbose"])
         self.assertEqual(plan["output_file"], "myfile.ipynb")
+        with self.assertRaises(IOError):
+            parse_plan(["this-file-doesn't-exist"])
 
     def test_annotate_source_path(self):
         nb_path = os.path.join(FIXTURES_DIR, "1_Intro.ipynb")
@@ -48,16 +70,16 @@ def test_annotate_source_path(self):
         self.assertEqual(nb.cells[0].metadata['xylophone'],
                          os.path.join('fixtures', '1_Intro.ipynb'))
 
-    def test_main(self):
-        if not hasattr(sys.stdout, "getvalue"):
-            self.fail("need to run in buffered mode")
+    def test_main_to_stdout(self):
+        main(TARGET_NBS)
+        self._validate_merged_three(reads(sys.stdout.getvalue(), as_version=4))
 
-        prior_args = sys.argv
+    def test_main_to_file(self):
+        _, path = mkstemp()
 
         try:
-            sys.argv = ['nbmerge'] + TARGET_NBS
-            main()
-        finally:
-            sys.argv = prior_args
-
-        self._validate_merged_three(reads(sys.stdout.getvalue(), as_version=4))
+            main(TARGET_NBS + ["-o", path])
+            with open(path) as fp:
+                self._validate_merged_three(reads(fp.read(), as_version=4))
+        except:
+            os.unlink(path)