✨ Basic download handling

harambe/core.py

@@ -1,4 +1,5 @@
 import asyncio
+import tempfile
 import uuid
 from functools import wraps
 from typing import Callable, List, Optional, Protocol, Union, Awaitable
@@ -15,7 +16,12 @@
     ResourceType,
     UnnecessaryResourceHandler,
 )
-from harambe.observer import LocalStorageObserver, LoggingObserver, OutputObserver
+from harambe.observer import (
+    LocalStorageObserver,
+    LoggingObserver,
+    OutputObserver,
+    DownloadMeta,
+)
 from harambe.tracker import FileDataTracker
 from harambe.types import URL, AsyncScraperType, Context, ScrapeResult, Stage
 
@@ -148,6 +154,47 @@ async def capture_url(
 
         return handler.captured_url()
 
+    async def capture_download(
+        self,
+        clickable: ElementHandle,
+    ) -> DownloadMeta:
+        """
+        Capture the download of a click event. This will click the element, download the resulting file
+        and apply some download handling logic from the observer to transform to a usable URL
+        """
+
+        async with self.page.expect_download() as download_info:
+            await clickable.click()
+        download = await download_info.value
+
+        # Create a temporary file to save the download
+        with tempfile.NamedTemporaryFile() as temp_file:
+            await download.save_as(temp_file.name)
+            with open(temp_file.name, "rb") as f:
+                content = f.read()
+
+        res = await asyncio.gather(
+            *[
+                o.on_download(download.suggested_filename, content)
+                for o in self._observers
+            ]
+        )
+        return res[0]
+
+    async def capture_pdf(
+        self,
+    ) -> DownloadMeta:
+        """
+        Capture the current page as a pdf and then apply some download handling logic
+        from the observer to transform to a usable URL
+        """
+        pdf_content = await self.page.pdf()
+        file_name = f"{self.page.url}-screen.pdf"
+        res = await asyncio.gather(
+            *[o.on_download(file_name, pdf_content) for o in self._observers]
+        )
+        return res[0]
+
     @staticmethod
     async def run(
         scraper: AsyncScraperType,
@@ -174,7 +221,7 @@ async def run(
                 viewport={"width": 1280, "height": 1024},
                 ignore_https_errors=True,
                 user_agent="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36"
-                           " (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36",
+                " (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36",
             )
             ctx.set_default_timeout(60000)
 
@@ -199,6 +246,7 @@ async def run(
                     context,
                 )
             except Exception as e:
+                # TODO: Fix path for non Mr. Watkins
                 await ctx.tracing.stop(
                     path="/Users/awtkns/PycharmProjects/harambe-public/trace.zip"
                 )

harambe/observer.py

@@ -1,17 +1,24 @@
-from typing import Any, Dict, List, Protocol, Tuple, runtime_checkable
+from abc import abstractmethod
+from typing import Any, Dict, List, Protocol, Tuple, runtime_checkable, TypedDict
 
 from harambe.tracker import FileDataTracker
 from harambe.types import URL, Context, Stage
 
 
 @runtime_checkable
 class OutputObserver(Protocol):
+    @abstractmethod
     async def on_save_data(self, data: Dict[str, Any]):
         raise NotImplementedError()
 
+    @abstractmethod
     async def on_queue_url(self, url: URL, context: Dict[str, Any]) -> None:
         raise NotImplementedError()
 
+    @abstractmethod
+    async def on_download(self, filename: str, content: bytes) -> "DownloadMeta":
+        raise NotImplementedError()
+
 
 class LoggingObserver(OutputObserver):
     # TODO: use logger
@@ -21,6 +28,13 @@ async def on_save_data(self, data: Dict[str, Any]):
     async def on_queue_url(self, url: URL, context: Dict[str, Any]) -> None:
         print(f"Enqueuing: {url} with context {context}")
 
+    async def on_download(self, filename: str, content: bytes) -> "DownloadMeta":
+        print(f"Downloading file: {filename}")  # TODO: use logger
+        return {
+            "url": f"www.test.com/download/{filename}",
+            "filename": filename,
+        }
+
 
 class LocalStorageObserver(OutputObserver):
     def __init__(self, domain: str, stage: Stage):
@@ -32,22 +46,48 @@ async def on_save_data(self, data: Dict[str, Any]):
     async def on_queue_url(self, url: URL, context: Dict[str, Any]) -> None:
         self._tracker.save_data({"url": url, "context": context})
 
+    async def on_download(self, filename: str, content: bytes) -> "DownloadMeta":
+        data = {
+            "url": f"www.test.com/download/{filename}",
+            "filename": filename,
+        }
+        self._tracker.save_data(data)
+        return data
+
 
 class InMemoryObserver(OutputObserver):
     def __init__(self):
-        self._urls: List[Tuple[URL, Context]] = []
         self._data: List[Dict[str, Any]] = []
+        self._urls: List[Tuple[URL, Context]] = []
+        self._files: List[Tuple[str, bytes]] = []
 
     async def on_save_data(self, data: Dict[str, Any]):
         self._data.append(data)
 
     async def on_queue_url(self, url: URL, context: Dict[str, Any]) -> None:
         self._urls.append((url, context))
 
+    async def on_download(self, filename: str, content: bytes) -> "DownloadMeta":
+        data = {
+            "url": f"www.test.com/download/{filename}",
+            "filename": filename,
+        }
+        self._files.append((filename, content))
+        return data
+
     @property
     def data(self) -> List[Dict[str, Any]]:
         return self._data
 
     @property
     def urls(self) -> List[Tuple[URL, Context]]:
         return self._urls
+
+    @property
+    def files(self) -> List[Tuple[str, bytes]]:
+        return self._files
+
+
+class DownloadMeta(TypedDict):
+    url: str
+    filename: str