Utility to check memory usage of an object graph

Author: jsbueno

README.md

@@ -198,4 +198,34 @@ Write out each step as a simple function (even a lambda), sync or async as
 is simpler, setup the Pipeline, and make a single call to have
 everything done.
 
+```python
+
+from extraasync import Pipeline
+import httpx
+
+urls = ["https://example.com/data1", "https://example.com/data2", ...] # (X 100s) 
+
+async def fetch(url):
+    client = httpx.AsyncClient()
+    return (await client.get(url)).text
+
+async def extract(html):
+    return re.findall(r"\<h1\>(.+)?\<\/h1\>", html)
+
+async def record(msg):
+    with open("output.txt", "at") as f:
+        f.write(msg + "\n")
+        
 
+async def main():
+    await Pipeline(urls, fetch, extract, record, max_concurrency=30, rate_limit=10)
+
+
+# Or, the fancy style:
+
+from extraasync.pipeline import process
+
+async def fancy_main():
+    await (urls >> Pipeline(rate_limit=10) | fetch | extract | record)
+    
+```

extraasync/__init__.py

@@ -7,4 +7,5 @@
 __version__ = "0.3.0"
 
 
-__all__ = ["aenumerate", "ExtraTaskGroup", "sync_to_async", "async_to_sync", "at_loop_stop_callback", "Pipeline", "remove_loop_stop_callback", "RateLimiter", "__version__"]
+__all__ = ["aenumerate", "ExtraTaskGroup", "sync_to_async", "async_to_sync",
+           "at_loop_stop_callback", "Pipeline", "remove_loop_stop_callback", "RateLimiter", "__version__"]

extraasync/aenumerate.py

@@ -1,4 +1,10 @@
 class aenumerate:
+    """
+    Asynchronous version of enumerate:
+
+    can be used in async-for loops
+    """
+
     def __init__(self, iterable, start=None):
         self.iterable = iterable
         self.start = start or 0
@@ -19,4 +25,3 @@ async def __aiter__(self):
 
     def __repr__(self):
         return f"Asynchronous enumeration for {self.iterable!r} at index {self.index!r}"
-

extraasync/pipeline.py

@@ -76,7 +76,6 @@ class AutoSet(MutableSet):
     Whenever an item is removed/discarded, an item is fetched from the queue as
     a task factory-no argument callable: it is called and added to the set.
 
-
     """
 
     def __init__(self, initial=()):
@@ -189,7 +188,8 @@ def normalized(self):
                 value /= 24 * 3600
             case _:
                 raise ValueError(
-                    f"Invalid time unit for frequency throtle - should be one of {TIME_UNIT}"
+                    f"Invalid time unit for frequency throtle - should be one of {
+                        TIME_UNIT}"
                 )
         return 1 / value
 
@@ -282,7 +282,8 @@ def put(self, value: tuple[int, t.Any]):
         if self.max_concurrency in (None, 0) or len(self.tasks) < self.max_concurrency:
             self.tasks.add(self._create_task(value))
         else:
-            self.tasks.queue.put_nowait(lambda value=value: self._create_task(value))
+            self.tasks.queue.put_nowait(
+                lambda value=value: self._create_task(value))
 
     def __call__(self, value):
         "just run the stage"
@@ -342,7 +343,8 @@ def __init__(
         """
         self.max_concurrency = max_concurrency
         self.data = (
-            _as_async_iterable(source) if source not in (None, Placeholder) else None
+            _as_async_iterable(source) if source not in (
+                None, Placeholder) else None
         )
         self.preserve_order = preserve_order
         # TBD: maybe allow limitting total memory usage instead of elements in the pipeline?
@@ -380,7 +382,8 @@ def _create_stages(self, stages):
 
     def reset(self):
         self.ordered_results = Heap()
-        self.output: asyncio.Queue[tuple[int | EXC_MARKER], t.Any] = asyncio.Queue()
+        self.output: asyncio.Queue[tuple[int |
+                                         EXC_MARKER], t.Any] = asyncio.Queue()
         self._create_stages(self.raw_stages)
         self.count = 0
 
@@ -435,7 +438,8 @@ async def __aiter__(self):
                     elif self.on_error == "strict":
                         raise result_data[1]
                     elif self.on_error == "lazy":
-                        raise NotImplementedError("Lazy error raising in pipeline")
+                        raise NotImplementedError(
+                            "Lazy error raising in pipeline")
                 if not self.preserve_order:
                     yield result_data
                 else:

extraasync/utils.py

@@ -0,0 +1,90 @@
+import asyncio
+import sys
+from array import array
+from collections.abc import Callable
+from types import FrameType, ModuleType, CodeType, CoroutineType
+from datetime import datetime, date, timedelta
+
+
+FrameLocalsProxy = (lambda: type(sys._getframe().f_locals))()
+
+DontRecurseTypes = float, int, complex, bool, datetime, date, timedelta, str, bytes, bytearray, array, FrameType
+ZeroCheckTypes = ModuleType, Callable, type, CodeType
+InspectAllAttrs = asyncio.Future, CoroutineType
+
+
+def sizeof(obj, seen=None):
+    """Recursively finds the aproximate size of a Python object graph,
+    including all references to non-class objects, and data which might be in
+    co-routine closures.
+
+    This is not perfect.
+    Nothing of this type will ever be perfect - recursivelly trying to find
+    all attributes referenced by an object will usually get out of hand.
+    Maybe use the GC to try to filter objects referenced only in the requested graph?
+
+    Nonetheless this implementation will include the size of objects currently held
+    in local variables inside asyncio tasks. For the typical workload here, it will sun up
+    the sizes of eventual API or HTTP request responses that are in temporary
+    use in pipeline modules
+    """
+    if seen is None:
+        seen = set()
+        try:
+            loop = asyncio.get_running_loop()
+            seen.add(id(loop))
+        except RuntimeError:
+            pass
+    if id(obj) in seen:
+        return 0
+    size = sys.getsizeof(obj)
+    seen.add(id(obj))
+    if isinstance(obj, ZeroCheckTypes):
+        # we just want isntance data
+        return 0
+
+    if isinstance(obj, FrameType):
+        if obj.f_locals is not obj.f_globals:
+            size += sizeof(obj.f_locals, None)
+
+    if isinstance(obj, DontRecurseTypes):
+        return size
+
+    if isinstance(obj, dict):
+        # for userdeifned Sequences, Mappings, etc...we get the contents size by attribute introspection.
+        for key, value in obj.items():
+            size += sizeof(key, seen)
+            size += sizeof(value, seen)
+    elif isinstance(obj, FrameLocalsProxy):
+        for value in obj.values():
+            size += sizeof(value, seen)
+    elif hasattr(type(obj), "__len__") and not isinstance(obj, range):
+        # take sequences. Skip generators. If seuqnece items are generated, we incorrectly add their sizes.
+        for item in obj:
+            size += sizeof(item, seen)
+    instance_checked = False
+    if hasattr(obj, "__dict__"):
+        size += sizeof(obj.__dict__, seen)
+        instance_checked = True
+    if hasattr(obj, "__slots__"):
+        for slot in obj.__slots__:
+            size += sizeof(getattr(obj, slot), seen)
+        instance_checked = True
+    if (frame := getattr(obj, "cr_frame", None)) or (frame := getattr(obj, "gi_frame", None)):
+        if isinstance(frame, FrameType):
+            size += sizeof(frame.f_locals, seen)
+        instance_checked = True
+    if not instance_checked and isinstance(obj, InspectAllAttrs):
+        # Certain objects created in native code, such as asyncio.Tasks won't have neither a __dict__ nor __slots__
+        # nor will "vars()" work on them.
+        # DOUBT: maybe use an allow-list here, instead of introspecting everything?
+        cls = type(obj)
+        for attrname in dir(obj):
+            attr = getattr(obj, attrname)
+            clsattr = getattr(cls, attrname, None)
+            if attr is clsattr or callable(attr):
+                # skip methods and class attributes
+                continue
+            size += sizeof(attr, seen)
+
+    return size

tests/test_pipeline.py

@@ -544,12 +544,12 @@ async def test_pipeline_add_data_and_execute_l_rhift_operator(): ...
 
 @pytest.mark.skip
 @pytest.mark.asyncio
-async def test_pipeline_store_result_r_rshift_operator(): ...
+async def test_pipeline_fine_tune_stages(): ...
 
 
 @pytest.mark.skip
 @pytest.mark.asyncio
-async def test_pipeline_fine_tune_stages(): ...
+async def test_pipeline_retry_simple(): ...
 
 
 @pytest.mark.asyncio