python-documentcloud 4.5.0__tar.gz → 4.7.0__tar.gz

{python_documentcloud-4.5.0/python_documentcloud.egg-info → python_documentcloud-4.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: python-documentcloud
-Version: 4.5.0
+Version: 4.7.0
 Summary: A simple Python wrapper for the DocumentCloud API
 Home-page: https://github.com/muckrock/python-documentcloud
 Author: Mitchell Kotler
@@ -27,6 +27,7 @@ Requires-Dist: urllib3
 Requires-Dist: pyyaml
 Requires-Dist: fastjsonschema
 Requires-Dist: python-squarelet
+Requires-Dist: token-bucket
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: coverage; extra == "dev"
@@ -37,6 +38,7 @@ Requires-Dist: twine; extra == "dev"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
+Requires-Dist: pytest-xdist; extra == "test"
 Requires-Dist: pytest-recording; extra == "test"
 Requires-Dist: vcrpy; extra == "test"
 Dynamic: author

{python_documentcloud-4.5.0 → python_documentcloud-4.7.0}/documentcloud/addon.py

@@ -182,6 +182,26 @@ class AddOn(BaseAddOn):
             f"addon_runs/{self.id}/", json={"file_name": file_name}
         )
 
+    def load_run_data(self):
+        "Load persistent data from this run"
+        if not self.id:
+            return {}
+
+        response = self.client.get(f"addon_runs/{self.id}/")
+        response.raise_for_status()
+        return response.json().get("data", {})
+
+    def store_run_data(self, data):
+        "Store persistent data for this run"
+        if not self.id:
+            print("Run ID not set. Try again later or check if something went wrong.")
+            return None
+
+        if not isinstance(data, dict):
+            raise TypeError("Invalid data")
+
+        return self.client.patch(f"addon_runs/{self.id}/", json={"data": data})
+
     def load_event_data(self):
         """Load persistent data for this event"""
         if not self.event_id:

python_documentcloud-4.7.0/documentcloud/client.py

@@ -0,0 +1,100 @@
+# Standard Library
+import logging
+import time
+
+# Third Party
+import token_bucket
+from squarelet import SquareletClient
+
+# Local
+from .documents import DocumentClient
+from .organizations import OrganizationClient
+from .projects import ProjectClient
+from .users import UserClient
+
+logger = logging.getLogger("documentcloud")
+
+# Per-endpoint rate limits applied on top of the global squarelet limit.
+# Format: (method, url_pattern, rate_per_second, capacity)
+#
+# Endpoint              Rate        Burst   Notes
+# --------              ----        -----   -----
+# GET  documents/search  15/min      50
+# POST documents/        12/min      100     25 docs/bulk call = up to 300 docs/min
+# PUT  documents/        12/min      100     25 docs/bulk call = up to 300 docs/min
+# GET  files/            15/min      100     PDFs, full text, and other private assets
+ENDPOINT_RATE_LIMITS = [
+    ("GET", "documents/search", 15 / 60, 50),
+    ("POST", "documents/", 12 / 60, 100),
+    ("PUT", "documents/", 12 / 60, 100),
+    ("GET", "files/", 15 / 60, 100),
+]
+
+
+class DocumentCloud(SquareletClient):
+    """
+    The public interface for the DocumentCloud API, now integrated with SquareletClient
+    """
+
+    def __init__(
+        self,
+        username=None,
+        password=None,
+        base_uri="https://api.www.documentcloud.org/api/",
+        auth_uri="https://accounts.muckrock.com/api/",
+        timeout=20,
+        loglevel=None,
+        rate_limit=True,
+        rate_limit_sleep=True,
+    ):
+        # Initialize SquareletClient for authentication and request handling
+        super().__init__(
+            base_uri=base_uri,
+            username=username,
+            password=password,
+            auth_uri=auth_uri,
+            timeout=timeout,
+            rate_limit=rate_limit,
+            rate_limit_sleep=rate_limit_sleep,
+        )
+
+        # Set up logging
+        if loglevel:
+            logging.basicConfig(
+                level=loglevel,
+                format="%(asctime)s %(levelname)-8s %(name)-25s %(message)s",
+            )
+        else:
+            logger.addHandler(logging.NullHandler())
+
+        # Build per-endpoint token bucket rate limiters
+        storage = token_bucket.MemoryStorage()
+        self._endpoint_limiters = [
+            (
+                pattern_method,
+                pattern,
+                token_bucket.Limiter(rate=rate, capacity=capacity, storage=storage),
+                f"{pattern_method}:{pattern}",
+            )
+            for pattern_method, pattern, rate, capacity in ENDPOINT_RATE_LIMITS
+        ]
+
+        # Initialize the sub-clients using SquareletClient
+        self.documents = DocumentClient(self)
+        self.projects = ProjectClient(self)
+        self.users = UserClient(self)
+        self.organizations = OrganizationClient(self)
+
+    def request(self, method, url, raise_error=True, **kwargs):
+        for pattern_method, pattern, limiter, bucket_key in self._endpoint_limiters:
+            if pattern_method.upper() == method.upper() and pattern in url:
+                if not limiter.consume(bucket_key):
+                    logger.warning(
+                        "Rate limit reached for %s %s, throttling...",
+                        method.upper(),
+                        pattern,
+                    )
+                    while not limiter.consume(bucket_key):
+                        time.sleep(0.1)
+                return super().request(method, url, raise_error=raise_error, **kwargs)
+        return super().request(method, url, raise_error=raise_error, **kwargs)

{python_documentcloud-4.5.0 → python_documentcloud-4.7.0}/documentcloud/documents.py

@@ -7,10 +7,13 @@ import datetime
 import logging
 import os
 import re
+import time
 import warnings
 from functools import partial
+from urllib.parse import urlparse
 
 # Third Party
+import token_bucket
 from requests.exceptions import RequestException
 
 # Local
@@ -23,15 +26,12 @@ from .sections import SectionClient
 from .toolbox import grouper, is_url, merge_dicts, requests_retry_session
 from .users import User
 
-try:
-    from urllib.parse import urlparse
-except ImportError:
-    from urlparse import urlparse
-
 logger = logging.getLogger("documentcloud")
 
 IMAGE_SIZES = ["thumbnail", "small", "normal", "large", "xlarge"]
 
+DEFAULT_USER_AGENT = "python-documentcloud"
+
 
 class Document(BaseAPIObject):
     """A single DocumentCloud document"""
@@ -168,12 +168,17 @@ class Document(BaseAPIObject):
 
         if base_netloc == url_netloc:
             # if the url host is the same as the base api host,
-            # sent the request with the client in order to include
+            # send the request with the client in order to include
             # authentication credentials
             response = self._client.get(url, full_url=True)
         else:
-            response = requests_retry_session().get(
-                url, headers={"User-Agent": "python-documentcloud2"}
+            response = self._client.documents.asset_get(
+                url,
+                headers={
+                    "User-Agent": self._client.session.headers.get(
+                        "User-Agent", DEFAULT_USER_AGENT
+                    )
+                },
             )
         if fmt == "text":
             return response.content.decode("utf8")
@@ -250,6 +255,26 @@ class DocumentClient(BaseAPIClient):
     api_path = "documents"
     resource = Document
 
+    def __init__(self, client):
+        super().__init__(client)
+        # Rate limit for public document asset fetches (S3-hosted).
+        # Private document assets go through the API client and are limited there.
+        # Token bucket: burst of 100, sustained at 15/min (0.25/sec).
+        storage = token_bucket.MemoryStorage()
+        self._asset_limiter = token_bucket.Limiter(
+            rate=15 / 60,
+            capacity=100,
+            storage=storage,
+        )
+        self._asset_session = requests_retry_session()
+
+    def asset_get(self, url, **kwargs):
+        if not self._asset_limiter.consume("asset"):
+            logger.warning("Rate limit reached for asset fetch, throttling...")
+            while not self._asset_limiter.consume("asset"):
+                time.sleep(0.1)
+        return self._asset_session.get(url, **kwargs)
+
     def search(self, query, **params):
         """Return documents matching a search query"""
 

python_documentcloud-4.7.0/documentcloud/exceptions.py

@@ -0,0 +1,15 @@
+"""
+Custom exceptions for python-documentcloud
+"""
+
+# Third Party
+# pylint: disable=unused-import
+# Import exceptions from python-squarelet
+from squarelet.exceptions import (
+    APIError,
+    CredentialsFailedError,
+    DoesNotExistError,
+    DuplicateObjectError,
+    MultipleObjectsReturnedError,
+    SquareletError as DocumentCloudError,
+)

{python_documentcloud-4.5.0 → python_documentcloud-4.7.0/python_documentcloud.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: python-documentcloud
-Version: 4.5.0
+Version: 4.7.0
 Summary: A simple Python wrapper for the DocumentCloud API
 Home-page: https://github.com/muckrock/python-documentcloud
 Author: Mitchell Kotler
@@ -27,6 +27,7 @@ Requires-Dist: urllib3
 Requires-Dist: pyyaml
 Requires-Dist: fastjsonschema
 Requires-Dist: python-squarelet
+Requires-Dist: token-bucket
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: coverage; extra == "dev"
@@ -37,6 +38,7 @@ Requires-Dist: twine; extra == "dev"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
+Requires-Dist: pytest-xdist; extra == "test"
 Requires-Dist: pytest-recording; extra == "test"
 Requires-Dist: vcrpy; extra == "test"
 Dynamic: author

{python_documentcloud-4.5.0 → python_documentcloud-4.7.0}/python_documentcloud.egg-info/SOURCES.txt

@@ -20,6 +20,7 @@ python_documentcloud.egg-info/SOURCES.txt
 python_documentcloud.egg-info/dependency_links.txt
 python_documentcloud.egg-info/requires.txt
 python_documentcloud.egg-info/top_level.txt
+tests/test_addon.py
 tests/test_annotations.py
 tests/test_base.py
 tests/test_client.py

{python_documentcloud-4.5.0 → python_documentcloud-4.7.0}/python_documentcloud.egg-info/requires.txt

@@ -7,6 +7,7 @@ urllib3
 pyyaml
 fastjsonschema
 python-squarelet
+token-bucket
 
 [dev]
 black
@@ -19,5 +20,6 @@ twine
 [test]
 pytest
 pytest-mock
+pytest-xdist
 pytest-recording
 vcrpy

{python_documentcloud-4.5.0 → python_documentcloud-4.7.0}/setup.py

@@ -7,7 +7,7 @@ with open("README.md", "r") as fh:
 
 setup(
     name="python-documentcloud",
-    version="4.5.0",
+    version="4.7.0",
     description="A simple Python wrapper for the DocumentCloud API",
     author="Mitchell Kotler",
     author_email="mitch@muckrock.com",
@@ -27,6 +27,7 @@ setup(
         "pyyaml",
         "fastjsonschema",
         "python-squarelet",
+        "token-bucket",
     ),
     extras_require={
         "dev": [
@@ -40,6 +41,7 @@ setup(
         "test": [
             "pytest",
             "pytest-mock",
+            "pytest-xdist",
             "pytest-recording",
             "vcrpy",
         ],

python_documentcloud-4.7.0/tests/test_addon.py

@@ -0,0 +1,141 @@
+# Standard Library
+from unittest.mock import MagicMock
+
+# Third Party
+import pytest
+
+# DocumentCloud
+from documentcloud.addon import AddOn
+
+# pylint: disable=redefined-outer-name
+
+
+@pytest.fixture
+def addon():
+    """An AddOn instance built without invoking argparse or constructing a real client.
+
+    Tests can override `.id`, `.event_id`, `.client`, etc. as needed.
+    """
+    instance = AddOn.__new__(AddOn)
+    instance.id = "run-123"
+    instance.addon_id = "addon-1"
+    instance.event_id = None
+    instance.documents = None
+    instance.query = None
+    instance.user_id = None
+    instance.org_id = None
+    instance.data = {}
+    instance.title = "Test AddOn"
+    instance.client = MagicMock()
+    return instance
+
+
+class TestLoadRunData:
+    def test_returns_data_when_run_id_set(self, addon):
+        addon.client.get.return_value.json.return_value = {"data": {"foo": "bar"}}
+
+        result = addon.load_run_data()
+
+        addon.client.get.assert_called_once_with("addon_runs/run-123/")
+        assert result == {"foo": "bar"}
+
+    def test_returns_empty_dict_when_no_run_id(self, addon):
+        addon.id = None
+
+        assert addon.load_run_data() == {}
+        addon.client.get.assert_not_called()
+
+    def test_returns_empty_dict_when_data_missing_from_response(self, addon):
+        addon.client.get.return_value.json.return_value = {}
+
+        assert addon.load_run_data() == {}
+
+
+class TestStoreRunData:
+    def test_patches_run_with_data(self, addon):
+        addon.store_run_data({"foo": "bar"})
+
+        addon.client.patch.assert_called_once_with(
+            "addon_runs/run-123/", json={"data": {"foo": "bar"}}
+        )
+
+    def test_no_op_when_no_run_id(self, addon, capsys):
+        addon.id = None
+
+        result = addon.store_run_data({"foo": "bar"})
+
+        assert result is None
+        addon.client.patch.assert_not_called()
+        assert "Run ID not set" in capsys.readouterr().out
+
+    def test_rejects_non_dict_data(self, addon):
+        with pytest.raises(TypeError):
+            addon.store_run_data("not a dict")
+
+        addon.client.patch.assert_not_called()
+
+
+class TestLoadEventData:
+    def test_returns_scratch_when_event_id_set(self, addon):
+        addon.event_id = "evt-9"
+        addon.client.get.return_value.json.return_value = {"scratch": {"x": 1}}
+
+        result = addon.load_event_data()
+
+        addon.client.get.assert_called_once_with("addon_events/evt-9/")
+        assert result == {"x": 1}
+
+    def test_returns_none_when_no_event_id(self, addon):
+        assert addon.load_event_data() is None
+        addon.client.get.assert_not_called()
+
+
+class TestStoreEventData:
+    def test_patches_event_with_scratch(self, addon):
+        addon.event_id = "evt-9"
+
+        addon.store_event_data({"x": 1})
+
+        addon.client.patch.assert_called_once_with(
+            "addon_events/evt-9/", json={"scratch": {"x": 1}}
+        )
+
+    def test_no_op_when_no_event_id(self, addon):
+        assert addon.store_event_data({"x": 1}) is None
+        addon.client.patch.assert_not_called()
+
+
+@pytest.fixture
+def real_addon(client, addon_run):
+    """An AddOn wired to the real `client` fixture and a freshly created run."""
+    instance = AddOn.__new__(AddOn)
+    instance.id = addon_run
+    instance.addon_id = None
+    instance.event_id = None
+    instance.documents = None
+    instance.query = None
+    instance.user_id = None
+    instance.org_id = None
+    instance.data = {}
+    instance.title = "Test AddOn"
+    instance.client = client
+    return instance
+
+
+class TestRunDataVCR:
+    """VCR-recorded round-trip tests against the dev DC.
+
+    Recording: set DC_TEST_ADDON_RUN_ID to an existing AddOnRun UUID on your
+    local dev DC, then run `make test-dev` (or `pytest --record-mode=new_episodes`).
+    """
+
+    def test_load_run_data_returns_dict(self, real_addon):
+        result = real_addon.load_run_data()
+        assert isinstance(result, dict)
+
+    def test_store_then_load_run_data_round_trip(self, real_addon):
+        payload = {"foo": "bar", "n": 42}
+        real_addon.store_run_data(payload)
+        loaded = real_addon.load_run_data()
+        assert loaded.get("foo") == "bar"
+        assert loaded.get("n") == 42

{python_documentcloud-4.5.0 → python_documentcloud-4.7.0}/tests/test_client.py

@@ -9,6 +9,7 @@ import pytest
 import ratelimit
 
 # DocumentCloud
+from documentcloud import DocumentCloud
 from documentcloud.constants import RATE_LIMIT
 from documentcloud.exceptions import APIError, CredentialsFailedError
 
@@ -111,3 +112,77 @@ def test_expired_refresh_token(short_client, record_mode):
     assert short_client.users.get("me")
     # check the refresh token was updated
     assert old_refresh_token != short_client.refresh_token
+
+
+def test_endpoint_rate_limit_burst_exhaustion():
+    """Token bucket should block after burst capacity is exhausted"""
+    client = DocumentCloud()
+    # Exhaust the search burst (capacity=50)
+    _pattern_method, _pattern, limiter, bucket_key = client._endpoint_limiters[0]
+    for _ in range(50):
+        limiter.consume(bucket_key)
+    assert not limiter.consume(bucket_key)
+
+
+def test_endpoint_rate_limit_method_specificity():
+    """GET and POST to documents/ should use different limiters"""
+    client = DocumentCloud()
+    limiters = {(pm, p): lim for pm, p, lim, _ in client._endpoint_limiters}
+    assert limiters[("GET", "files/")] is not limiters[("POST", "documents/")]
+
+
+def test_endpoint_rate_limit_pattern_ordering():
+    """documents/search should match before documents/"""
+    client = DocumentCloud()
+    url = "documents/search/"
+    matched = next(
+        p for pm, p, _, _ in client._endpoint_limiters if pm == "GET" and p in url
+    )
+    assert matched == "documents/search"
+
+
+def test_asset_rate_limit_burst_exhaustion():
+    """Asset token bucket should block after burst capacity is exhausted"""
+    client = DocumentCloud()
+    limiter = client.documents._asset_limiter
+    for _ in range(100):
+        limiter.consume("asset")
+    assert not limiter.consume("asset")
+
+
+def test_asset_rate_limit_refills():
+    """Asset token bucket should refill over time"""
+    client = DocumentCloud()
+    limiter = client.documents._asset_limiter
+    for _ in range(100):
+        limiter.consume("asset")
+    assert not limiter.consume("asset")
+    time.sleep(5)
+    assert limiter.consume("asset")
+
+
+def test_endpoint_rate_limit_buckets_are_independent():
+    """Exhausting one endpoint's bucket should not affect another"""
+    client = DocumentCloud()
+    limiters = {(pm, p): (lim, bk) for pm, p, lim, bk in client._endpoint_limiters}
+    search_limiter, search_key = limiters[("GET", "documents/search")]
+    files_limiter, files_key = limiters[("GET", "files/")]
+
+    # Exhaust search bucket
+    for _ in range(50):
+        search_limiter.consume(search_key)
+    assert not search_limiter.consume(search_key)
+
+    # Files bucket should still have tokens
+    assert files_limiter.consume(files_key)
+
+
+def test_endpoint_rate_limit_no_match_for_unrecognized_url():
+    """Unrecognized URLs should not match any endpoint limiter"""
+    client = DocumentCloud()
+    url = "users/me/"
+    matched = next(
+        (p for pm, p, _, _ in client._endpoint_limiters if p in url),
+        None,
+    )
+    assert matched is None

python_documentcloud-4.5.0/documentcloud/client.py

@@ -1,58 +0,0 @@
-# Import SquareletClient from python-squarelet
-# Standard Library
-import logging
-
-# Third Party
-from squarelet import SquareletClient
-
-# Local
-# Local Imports
-from .documents import DocumentClient
-from .organizations import OrganizationClient
-from .projects import ProjectClient
-from .users import UserClient
-
-logger = logging.getLogger("documentcloud")
-
-
-class DocumentCloud(SquareletClient):
-    """
-    The public interface for the DocumentCloud API, now integrated with SquareletClient
-    """
-
-    def __init__(
-        self,
-        username=None,
-        password=None,
-        base_uri="https://api.www.documentcloud.org/api/",
-        auth_uri="https://accounts.muckrock.com/api/",
-        timeout=20,
-        loglevel=None,
-        rate_limit=True,
-        rate_limit_sleep=True,
-    ):
-        # Initialize SquareletClient for authentication and request handling
-        super().__init__(
-            base_uri=base_uri,
-            username=username,
-            password=password,
-            auth_uri=auth_uri,
-            timeout=timeout,
-            rate_limit=rate_limit,
-            rate_limit_sleep=rate_limit_sleep,
-        )
-
-        # Set up logging
-        if loglevel:
-            logging.basicConfig(
-                level=loglevel,
-                format="%(asctime)s %(levelname)-8s %(name)-25s %(message)s",
-            )
-        else:
-            logger.addHandler(logging.NullHandler())
-
-        # Initialize the sub-clients using SquareletClient
-        self.documents = DocumentClient(self)
-        self.projects = ProjectClient(self)
-        self.users = UserClient(self)
-        self.organizations = OrganizationClient(self)

python_documentcloud-4.5.0/documentcloud/exceptions.py

@@ -1,12 +0,0 @@
-"""
-Custom exceptions for python-documentcloud
-"""
-
-# pylint: disable=unused-import
-# Import exceptions from python-squarelet
-from squarelet.exceptions import SquareletError as DocumentCloudError
-from squarelet.exceptions import DuplicateObjectError
-from squarelet.exceptions import CredentialsFailedError
-from squarelet.exceptions import APIError
-from squarelet.exceptions import DoesNotExistError
-from squarelet.exceptions import MultipleObjectsReturnedError