lmnr 0.6.16__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/sdk/browser/playwright_otel.py

@@ -2,10 +2,14 @@ import logging
 import uuid
 
 from lmnr.opentelemetry_lib.utils.package_check import is_package_installed
-from lmnr.sdk.browser.pw_utils import handle_navigation_async, handle_navigation_sync
+from lmnr.sdk.browser.pw_utils import (
+    start_recording_events_async,
+    start_recording_events_sync,
+    take_full_snapshot,
+    take_full_snapshot_async,
+)
 from lmnr.sdk.browser.utils import with_tracer_and_client_wrapper
 from lmnr.sdk.client.asynchronous.async_client import AsyncLaminarClient
-from lmnr.sdk.client.synchronous.sync_client import LaminarClient
 from lmnr.version import __version__
 
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
@@ -13,29 +17,22 @@ from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.trace import (
     get_tracer,
     Tracer,
-    get_current_span,
-    Span,
-    INVALID_SPAN,
-    set_span_in_context,
 )
-from opentelemetry.context import get_current
 from typing import Collection
 from wrapt import wrap_function_wrapper
 
 try:
     if is_package_installed("playwright"):
-        from playwright.async_api import Browser, BrowserContext, Page
+        from playwright.async_api import Browser, BrowserContext
         from playwright.sync_api import (
             Browser as SyncBrowser,
             BrowserContext as SyncBrowserContext,
-            Page as SyncPage,
         )
     elif is_package_installed("patchright"):
-        from patchright.async_api import Browser, BrowserContext, Page
+        from patchright.async_api import Browser, BrowserContext
         from patchright.sync_api import (
             Browser as SyncBrowser,
             BrowserContext as SyncBrowserContext,
-            Page as SyncPage,
         )
     else:
         raise ImportError(
@@ -54,67 +51,26 @@ except ImportError as e:
 _instruments = ("playwright >= 1.9.0",)
 logger = logging.getLogger(__name__)
 
-_context_spans: dict[str, Span] = {}
-
-
-@with_tracer_and_client_wrapper
-def _wrap_new_page(
-    tracer: Tracer, client: LaminarClient, to_wrap, wrapped, instance, args, kwargs
-):
-    with tracer.start_as_current_span(
-        f"{to_wrap.get('object')}.{to_wrap.get('method')}"
-    ) as span:
-        page = wrapped(*args, **kwargs)
-        session_id = str(uuid.uuid4().hex)
-        trace_id = format(get_current_span().get_span_context().trace_id, "032x")
-        span.set_attribute("lmnr.internal.has_browser_session", True)
-        handle_navigation_sync(page, session_id, trace_id, client)
-        return page
-
-
-@with_tracer_and_client_wrapper
-async def _wrap_new_page_async(
-    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
-):
-    with tracer.start_as_current_span(
-        f"{to_wrap.get('object')}.{to_wrap.get('method')}"
-    ) as span:
-        page = await wrapped(*args, **kwargs)
-        session_id = str(uuid.uuid4().hex)
-        trace_id = format(span.get_span_context().trace_id, "032x")
-        span.set_attribute("lmnr.internal.has_browser_session", True)
-        await handle_navigation_async(page, session_id, trace_id, client)
-        return page
-
 
 @with_tracer_and_client_wrapper
 def _wrap_new_browser_sync(
-    tracer: Tracer, client: LaminarClient, to_wrap, wrapped, instance, args, kwargs
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
-    global _context_spans
     browser: SyncBrowser = wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
-    for context in browser.contexts:
-        span = get_current_span()
-        if span == INVALID_SPAN:
-            span = tracer.start_span(
-                name=f"{to_wrap.get('object')}.{to_wrap.get('method')}"
-            )
-            set_span_in_context(span, get_current())
-            _context_spans[id(context)] = span
-        span.set_attribute("lmnr.internal.has_browser_session", True)
-        trace_id = format(span.get_span_context().trace_id, "032x")
-
-        def handle_page_navigation(page: SyncPage):
-            return handle_navigation_sync(page, session_id, trace_id, client)
-
-        context.on(
-            "page",
-            handle_page_navigation,
-        )
 
+    def create_page_handler(session_id, client):
+        def page_handler(page):
+            start_recording_events_sync(page, session_id, client)
+
+        return page_handler
+
+    for context in browser.contexts:
+        page_handler = create_page_handler(session_id, client)
+        context.on("page", page_handler)
         for page in context.pages:
-            handle_navigation_sync(page, session_id, trace_id, client)
+            start_recording_events_sync(page, session_id, client)
+
     return browser
 
 
@@ -122,54 +78,41 @@ def _wrap_new_browser_sync(
 async def _wrap_new_browser_async(
     tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
-    global _context_spans
     browser: Browser = await wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
+
+    def create_page_handler(session_id, client):
+        async def page_handler(page):
+            await start_recording_events_async(page, session_id, client)
+
+        return page_handler
+
     for context in browser.contexts:
-        span = get_current_span()
-        if span == INVALID_SPAN:
-            span = tracer.start_span(
-                name=f"{to_wrap.get('object')}.{to_wrap.get('method')}"
-            )
-            set_span_in_context(span, get_current())
-            _context_spans[id(context)] = span
-        span.set_attribute("lmnr.internal.has_browser_session", True)
-        trace_id = format(span.get_span_context().trace_id, "032x")
-
-        async def handle_page_navigation(page: Page):
-            return await handle_navigation_async(page, session_id, trace_id, client)
-
-        context.on("page", handle_page_navigation)
+        page_handler = create_page_handler(session_id, client)
+        context.on("page", page_handler)
         for page in context.pages:
-            await handle_navigation_async(page, session_id, trace_id, client)
+            await start_recording_events_async(page, session_id, client)
     return browser
 
 
 @with_tracer_and_client_wrapper
 def _wrap_new_context_sync(
-    tracer: Tracer, client: LaminarClient, to_wrap, wrapped, instance, args, kwargs
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
     context: SyncBrowserContext = wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
-    span = get_current_span()
-    if span == INVALID_SPAN:
-        span = tracer.start_span(
-            name=f"{to_wrap.get('object')}.{to_wrap.get('method')}"
-        )
-        set_span_in_context(span, get_current())
-        _context_spans[id(context)] = span
-    span.set_attribute("lmnr.internal.has_browser_session", True)
-    trace_id = format(span.get_span_context().trace_id, "032x")
-
-    def handle_page_navigation(page: SyncPage):
-        return handle_navigation_sync(page, session_id, trace_id, client)
-
-    context.on(
-        "page",
-        handle_page_navigation,
-    )
+
+    def create_page_handler(session_id, client):
+        def page_handler(page):
+            start_recording_events_sync(page, session_id, client)
+
+        return page_handler
+
+    page_handler = create_page_handler(session_id, client)
+    context.on("page", page_handler)
     for page in context.pages:
-        handle_navigation_sync(page, session_id, trace_id, client)
+        start_recording_events_sync(page, session_id, client)
+
     return context
 
 
@@ -179,80 +122,58 @@ async def _wrap_new_context_async(
 ):
     context: BrowserContext = await wrapped(*args, **kwargs)
     session_id = str(uuid.uuid4().hex)
-    span = get_current_span()
-    if span == INVALID_SPAN:
-        span = tracer.start_span(
-            name=f"{to_wrap.get('object')}.{to_wrap.get('method')}"
-        )
-        set_span_in_context(span, get_current())
-        _context_spans[id(context)] = span
-    span.set_attribute("lmnr.internal.has_browser_session", True)
-    trace_id = format(span.get_span_context().trace_id, "032x")
 
-    async def handle_page_navigation(page):
-        return await handle_navigation_async(page, session_id, trace_id, client)
+    def create_page_handler(session_id, client):
+        async def page_handler(page):
+            await start_recording_events_async(page, session_id, client)
+
+        return page_handler
 
-    context.on("page", handle_page_navigation)
+    page_handler = create_page_handler(session_id, client)
+    context.on("page", page_handler)
     for page in context.pages:
-        await handle_navigation_async(page, session_id, trace_id, client)
+        await start_recording_events_async(page, session_id, client)
+
     return context
 
 
 @with_tracer_and_client_wrapper
-def _wrap_close_browser_sync(
-    tracer: Tracer,
-    client: LaminarClient,
-    to_wrap,
-    wrapped,
-    instance: SyncBrowser,
-    args,
-    kwargs,
+def _wrap_bring_to_front_sync(
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
+):
+    wrapped(*args, **kwargs)
+    take_full_snapshot(instance)
+
+
+@with_tracer_and_client_wrapper
+async def _wrap_bring_to_front_async(
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
-    global _context_spans
-    for context in instance.contexts:
-        key = id(context)
-        span = _context_spans.get(key)
-        if span:
-            if span.is_recording():
-                span.end()
-            _context_spans.pop(key)
-    return wrapped(*args, **kwargs)
+    await wrapped(*args, **kwargs)
+    await take_full_snapshot_async(instance)
 
 
 @with_tracer_and_client_wrapper
-async def _wrap_close_browser_async(
-    tracer: Tracer,
-    client: AsyncLaminarClient,
-    to_wrap,
-    wrapped,
-    instance: Browser,
-    args,
-    kwargs,
+def _wrap_browser_new_page_sync(
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
 ):
-    global _context_spans
-    for context in instance.contexts:
-        key = id(context)
-        span = _context_spans.get(key)
-        if span:
-            if span.is_recording():
-                span.end()
-            _context_spans.pop(key)
-    return await wrapped(*args, **kwargs)
+    page = wrapped(*args, **kwargs)
+    session_id = str(uuid.uuid4().hex)
+    start_recording_events_sync(page, session_id, client)
+    return page
+
+
+@with_tracer_and_client_wrapper
+async def _wrap_browser_new_page_async(
+    tracer: Tracer, client: AsyncLaminarClient, to_wrap, wrapped, instance, args, kwargs
+):
+    page = await wrapped(*args, **kwargs)
+    session_id = str(uuid.uuid4().hex)
+    await start_recording_events_async(page, session_id, client)
+    return page
 
 
 WRAPPED_METHODS = [
-    {
-        "package": "playwright.sync_api",
-        "object": "BrowserContext",
-        "method": "new_page",
-        "wrapper": _wrap_new_page,
-    },
-    {
-        "package": "playwright.sync_api",
-        "object": "Browser",
-        "method": "new_page",
-        "wrapper": _wrap_new_page,
-    },
     {
         "package": "playwright.sync_api",
         "object": "BrowserType",
@@ -271,12 +192,6 @@ WRAPPED_METHODS = [
         "method": "connect_over_cdp",
         "wrapper": _wrap_new_browser_sync,
     },
-    {
-        "package": "playwright.sync_api",
-        "object": "Browser",
-        "method": "close",
-        "wrapper": _wrap_close_browser_sync,
-    },
     {
         "package": "playwright.sync_api",
         "object": "Browser",
@@ -289,21 +204,21 @@ WRAPPED_METHODS = [
         "method": "launch_persistent_context",
         "wrapper": _wrap_new_context_sync,
     },
-]
-
-WRAPPED_METHODS_ASYNC = [
     {
-        "package": "playwright.async_api",
-        "object": "BrowserContext",
-        "method": "new_page",
-        "wrapper": _wrap_new_page_async,
+        "package": "playwright.sync_api",
+        "object": "Page",
+        "method": "bring_to_front",
+        "wrapper": _wrap_bring_to_front_sync,
     },
     {
-        "package": "playwright.async_api",
+        "package": "playwright.sync_api",
         "object": "Browser",
         "method": "new_page",
-        "wrapper": _wrap_new_page_async,
+        "wrapper": _wrap_browser_new_page_sync,
     },
+]
+
+WRAPPED_METHODS_ASYNC = [
     {
         "package": "playwright.async_api",
         "object": "BrowserType",
@@ -322,12 +237,6 @@ WRAPPED_METHODS_ASYNC = [
         "method": "connect_over_cdp",
         "wrapper": _wrap_new_browser_async,
     },
-    {
-        "package": "playwright.async_api",
-        "object": "Browser",
-        "method": "close",
-        "wrapper": _wrap_close_browser_async,
-    },
     {
         "package": "playwright.async_api",
         "object": "Browser",
@@ -340,13 +249,24 @@ WRAPPED_METHODS_ASYNC = [
         "method": "launch_persistent_context",
         "wrapper": _wrap_new_context_async,
     },
+    {
+        "package": "playwright.async_api",
+        "object": "Page",
+        "method": "bring_to_front",
+        "wrapper": _wrap_bring_to_front_async,
+    },
+    {
+        "package": "playwright.async_api",
+        "object": "Browser",
+        "method": "new_page",
+        "wrapper": _wrap_browser_new_page_async,
+    },
 ]
 
 
 class PlaywrightInstrumentor(BaseInstrumentor):
-    def __init__(self, client: LaminarClient, async_client: AsyncLaminarClient):
+    def __init__(self, async_client: AsyncLaminarClient):
         super().__init__()
-        self.client = client
         self.async_client = async_client
 
     def instrumentation_dependencies(self) -> Collection[str]:
@@ -356,6 +276,8 @@ class PlaywrightInstrumentor(BaseInstrumentor):
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, __version__, tracer_provider)
 
+        # Both sync and async methods use async_client because we are using
+        # a background asyncio loop for async sends
         for wrapped_method in WRAPPED_METHODS:
             wrap_package = wrapped_method.get("package")
             wrap_object = wrapped_method.get("object")
@@ -366,7 +288,7 @@ class PlaywrightInstrumentor(BaseInstrumentor):
                     f"{wrap_object}.{wrap_method}",
                     wrapped_method.get("wrapper")(
                         tracer,
-                        self.client,
+                        self.async_client,
                         wrapped_method,
                     ),
                 )
@@ -393,13 +315,8 @@ class PlaywrightInstrumentor(BaseInstrumentor):
 
     def _uninstrument(self, **kwargs):
         # Unwrap methods
-        global _context_spans
         for wrapped_method in WRAPPED_METHODS + WRAPPED_METHODS_ASYNC:
             wrap_package = wrapped_method.get("package")
             wrap_object = wrapped_method.get("object")
             wrap_method = wrapped_method.get("method")
             unwrap(wrap_package, f"{wrap_object}.{wrap_method}")
-            for span in _context_spans.values():
-                if span.is_recording():
-                    span.end()
-            _context_spans = {}