From dee1a2daaf13a18b59675c09813f3c10e0a05e13 Mon Sep 17 00:00:00 2001
From: Garrett Bischof <gbischof@bnl.gov>
Date: Mon, 10 Jun 2019 15:18:43 -0400
Subject: [PATCH] more changes from dan's suggestions

---
 intake_bluesky/core.py             | 21 +++++++++------------
 intake_bluesky/mongo_normalized.py |  4 ++--
 intake_bluesky/tests/test_core.py  |  2 ++
 3 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/intake_bluesky/core.py b/intake_bluesky/core.py
index d052c2e..9ea7408 100644
--- a/intake_bluesky/core.py
+++ b/intake_bluesky/core.py
@@ -19,7 +19,7 @@
 import heapq
 
 
-def to_event_pages(get_event_cursor):
+def to_event_pages(get_event_cursor, page_size):
     """
     Decorator that changes get_event_cursor get_event_pages.
 
@@ -33,11 +33,11 @@ def to_event_pages(get_event_cursor):
     -------
     get_event_pages : function
     """
-    @functools.wraps
+    @functools.wraps(get_event_cursor)
     def get_event_pages(*args, **kwargs):
         event_cursor = get_event_cursor(*args, **kwargs)
         while True:
-            result = list(itertools.islice(event_cursor, 2500))
+            result = list(itertools.islice(event_cursor, page_size))
             if result:
                 yield event_model.pack_event_page(*result)
             else:
@@ -45,7 +45,7 @@ def get_event_pages(*args, **kwargs):
     return get_event_pages
 
 
-def to_datum_pages(get_datum_cursor):
+def to_datum_pages(get_datum_cursor, page_size):
     """
     Decorator that changes get_datum_cursor get_datum_pages.
 
@@ -59,11 +59,11 @@ def to_datum_pages(get_datum_cursor):
     -------
     get_datum_pages : function
     """
-    @functools.wraps
+    @functools.wraps(get_datum_cursor)
     def get_datum_pages(*args, **kwargs):
         datum_cursor = get_datum_cursor(*args, **kwargs)
         while True:
-            result = list(itertools.islice(datum_cursor, 2500))
+            result = list(itertools.islice(datum_cursor, page_size))
             if result:
                 yield event_model.pack_datum_page(*result)
             else:
@@ -139,8 +139,7 @@ def documents_to_xarray(*, start_doc, stop_doc, descriptor_docs,
     get_resource : callable
         Expected signature ``get_resource(resource_uid) -> Resource``
     lookup_resource_for_datum : callable
-        Expected signature
-        ``lookup_resource_for_datum(datum_id) -> resource_uid``
+        Expected signature ``lookup_resource_for_datum(datum_id) -> resource_uid``
     get_datum_pages : callable
         Expected signature ``get_datum_pages(resource_uid) -> generator``
         where ``generator`` yields datum_page documents
@@ -420,8 +419,7 @@ class BlueskyRun(intake.catalog.Catalog):
     get_resource : callable
         Expected signature ``get_resource(resource_uid) -> Resource``
     lookup_resource_for_datum : callable
-        Expected signature
-        ``lookup_resource_for_datum(datum_id) -> resource_uid``
+        Expected signature ``lookup_resource_for_datum(datum_id) -> resource_uid``
     get_datum_pages : callable
         Expected signature ``get_datum_pages(resource_uid) -> generator``
         where ``generator`` yields Datum documents
@@ -626,8 +624,7 @@ class BlueskyEventStream(intake_xarray.base.DataSourceMixin):
     get_resource : callable
         Expected signature ``get_resource(resource_uid) -> Resource``
     lookup_resource_for_datum : callable
-        Expected signature
-        ``lookup_resource_for_datum(datum_id) -> resource_uid``
+        Expected signature ``lookup_resource_for_datum(datum_id) -> resource_uid``
     get_datum_pages : callable
         Expected signature ``get_datum_pages(resource_uid) -> generator``
         where ``generator`` yields datum_page documents
diff --git a/intake_bluesky/mongo_normalized.py b/intake_bluesky/mongo_normalized.py
index 2759d25..57b9a46 100644
--- a/intake_bluesky/mongo_normalized.py
+++ b/intake_bluesky/mongo_normalized.py
@@ -31,11 +31,11 @@ def get_run_start():
             get_run_start=get_run_start,
             get_run_stop=partial(self.catalog._get_run_stop, uid),
             get_event_descriptors=partial(self.catalog._get_event_descriptors, uid),
-            get_event_pages=to_event_pages(self.catalog._get_event_cursor),
+            get_event_pages=to_event_pages(self.catalog._get_event_cursor, 2500),
             get_event_count=self.catalog._get_event_count,
             get_resource=self.catalog._get_resource,
             lookup_resource_for_datum=self.catalog._lookup_resource_for_datum,
-            get_datum_pages=to_datum_pages(self.catalog._get_datum_cursor),
+            get_datum_pages=to_datum_pages(self.catalog._get_datum_cursor, 2500),
             filler=self.catalog.filler)
         return intake.catalog.local.LocalCatalogEntry(
             name=run_start_doc['uid'],
diff --git a/intake_bluesky/tests/test_core.py b/intake_bluesky/tests/test_core.py
index ac1760a..6cedc31 100644
--- a/intake_bluesky/tests/test_core.py
+++ b/intake_bluesky/tests/test_core.py
@@ -1,9 +1,11 @@
 import event_model
 from intake_bluesky.core import documents_to_xarray
 
+
 def no_event_pages(descriptor_uid):
     yield from ()
 
+
 def test_no_descriptors():
     run_bundle = event_model.compose_run()
     start_doc = run_bundle.start_doc