austin-python 1.7.1__py3-none-any.whl → 2.0.0__py3-none-any.whl

austin/__init__.py

@@ -20,310 +20,3 @@
 # GNU General Public License for more details.
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
-
-
-import argparse
-import functools
-import os
-import os.path
-from abc import ABC
-from abc import abstractmethod
-from itertools import takewhile
-from pathlib import Path
-from typing import Any
-from typing import Callable
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Tuple
-
-import psutil
-
-from austin.config import AustinConfiguration
-
-
-try:
-    _cached_property = functools.cached_property  # type: ignore[attr-defined]
-except AttributeError:
-
-    def _cached_property(f: Callable[..., Any]) -> property:  # type: ignore[no-redef]
-        return property(functools.lru_cache(maxsize=1)(f))
-
-
-SemVer = Tuple[int, int, int]
-
-
-def _to_semver(version: Optional[str]) -> SemVer:
-    if version is None:
-        return (0, 0, 0)
-
-    try:
-        return (  # type: ignore[return-value]
-            tuple(
-                int(_)
-                for _ in "".join(
-                    list(
-                        takewhile(
-                            lambda _: _.isdigit() or _ == ".", version.replace("?", "0")
-                        )
-                    )
-                ).split(".")
-            )
-            + (0, 0, 0)
-        )[:3]
-    except ValueError:
-        raise ValueError("Invalid semantic version") from None
-
-
-class AustinError(Exception):
-    """Basic Austin Error."""
-
-    pass
-
-
-class AustinTerminated(AustinError):
-    """Austin termination exception.
-
-    Thrown when Austin is terminated with a call to ``terminate``.
-    """
-
-    pass
-
-
-class BaseAustin(ABC):
-    """Base Austin class.
-
-    Defines the general API that abstract Austin as an external process.
-    Subclasses should implement the :func:`start` method and either define the
-    :func:`on_sample_received` method or pass it via the constructor.
-    Additionally, the :func:`on_ready` and the :func:`on_terminate` methods can
-    be overridden or passed via the constructor to catch the corresponding
-    events. Austin is considered to be ready when the first sample is received;
-    it is considered to have terminated when the process has terminated
-    gracefully.
-
-    If an error was encountered, the :class:`AustinError` exception is thrown.
-    """
-
-    BINARY = "austin" + (os.name == "nt" and ".exe" or "")
-    BINARY_VERSION = (3, 0, 0)
-
-    def __init__(
-        self,
-        sample_callback: Optional[Callable[[str], None]] = None,
-        ready_callback: Optional[
-            Callable[[psutil.Process, psutil.Process, str], None]
-        ] = None,
-        terminate_callback: Optional[Callable[[Dict[str, str]], None]] = None,
-    ) -> None:
-        """The ``BaseAustin`` constructor.
-
-        The callbacks can be passed as arguments, if they are not
-        defined/overridden in the definition of a sub-class.
-        """
-        self._proc: psutil.Process = None
-        self._child_proc: psutil.Process = None
-        self._cmd_line: Optional[str] = None
-        self._running: bool = False
-        self._meta: Dict[str, str] = {}
-
-        try:
-            self._sample_callback = sample_callback or self.on_sample_received  # type: ignore[attr-defined]
-        except AttributeError:
-            raise AustinError("No sample callback given or implemented.") from None
-
-        self._terminate_callback = terminate_callback or self.on_terminate
-        self._ready_callback = ready_callback or self.on_ready
-
-        self._version: Optional[str] = None
-        self._python_version: Optional[str] = None
-
-    def _get_process_info(
-        self, args: argparse.Namespace, austin_pid: int
-    ) -> Tuple[psutil.Process, psutil.Process, str]:
-        try:
-            self._proc = psutil.Process(austin_pid)
-        except psutil.NoSuchProcess:
-            raise AustinError("Cannot find Austin process.") from None
-
-        if not args.pid:  # Austin is forking
-            try:
-                self._child_proc = self._proc.children()[0]
-                if self._child_proc is None:
-                    raise IndexError
-            except IndexError:
-                raise AustinError("Cannot find Austin child process.") from None
-        else:  # Austin is attaching
-            try:
-                self._child_proc = psutil.Process(args.pid)
-            except psutil.NoSuchProcess:
-                raise AustinError(
-                    f"Cannot attach to process with invalid PID {args.pid}."
-                ) from None
-
-        self._cmd_line = " ".join(self._child_proc.cmdline())
-
-        self._running = True
-
-        return self._proc, self._child_proc, self._cmd_line
-
-    @abstractmethod
-    def start(self, args: Optional[List[str]] = None) -> Any:
-        """Start Austin.
-
-        Every subclass should implement this method and ensure that it spawns
-        a new Austin process.
-        """
-        ...
-
-    def get_process(self) -> psutil.Process:
-        """Get the underlying Austin process.
-
-        Return an instance of :class:`psuitl.Process` that can be used to
-        control the underlying Austin process at the OS level.
-        """
-        return self._proc
-
-    def get_command_line(self) -> Optional[str]:
-        """Get the inferred command line.
-
-        Return the command line of the (main) process that is being profiled
-        by Austin.
-        """
-        return self._cmd_line
-
-    def is_running(self) -> bool:
-        """Determine whether Austin is running."""
-        return self._running
-
-    def terminate(self, wait: bool = False) -> None:
-        """Terminate Austin.
-
-        Stop the underlying Austin process by sending a termination signal.
-        """
-        if not self._proc:
-            raise AustinError("Austin has not been started yet")
-
-        try:
-            self._proc.terminate()
-            if wait:
-                self._proc.wait()
-        except psutil.NoSuchProcess:
-            raise AustinError("Austin has already terminated") from None
-
-        self._running = False
-        self._proc = None
-        self._child_proc = None
-
-    def get_child_process(self) -> psutil.Process:
-        """Get the child process.
-
-        Return an instalce of :class:`psutil.Process` representing the (main)
-        process being profiled by Austin at the OS level.
-        """
-        return self._child_proc
-
-    def submit_sample(self, data: bytes) -> None:
-        """Submit a sample to the sample callback.
-
-        This method takes care of converting the raw binary data retrieved from
-        Austin into a Python string.
-        """
-        try:
-            sample = data.decode()
-        except UnicodeDecodeError:
-            try:
-                sample = data.decode("ascii")
-            except UnicodeDecodeError:
-                return
-
-        self._sample_callback(sample.rstrip())
-
-    def check_exit(self, rcode: int, stderr: Optional[str]) -> None:
-        """Check Austin exit status."""
-        if rcode:
-            if rcode in {-15, 15, 241}:
-                raise AustinTerminated()
-            raise AustinError(f"({rcode}) {stderr}")
-
-    def check_version(self) -> None:
-        """Check for the minimum Austin binary version."""
-        austin_version = self.version
-        if austin_version is None:
-            raise AustinError("Cannot determine Austin version")
-        if austin_version < self.BINARY_VERSION:
-            raise AustinError(
-                f"Incompatible Austin version (got {austin_version}, expected >= {self.BINARY_VERSION})"
-            )
-
-    # ---- Default callbacks ----
-
-    def on_terminate(self, stats: Dict[str, str]) -> Any:
-        """Terminate event callback.
-
-        Implement to be notified when Austin has terminated gracefully. The
-        callback accepts an argument that will receive the global statistics.
-        """
-        pass
-
-    def on_ready(
-        self,
-        process: psutil.Process,
-        child_process: psutil.Process,
-        command_line: str,
-        data: Any = None,
-    ) -> Any:
-        """Ready event callback.
-
-        Implement to get notified when Austin has successfully started or
-        attached the Python process to profile and the first sample has been
-        produced. This callback receives the Austin process and it's (main)
-        profiled process as instances of :class:`psutil.Process`, along with
-        the command line of the latter.
-        """
-        pass
-
-    # ---- Properties ----
-
-    @_cached_property
-    def binary_path(self) -> str:
-        """Discover the path of the Austin binary.
-
-        Lookup order is:
-        - current working directory
-        - ``AUSTINPATH`` variable
-        - ``~/.austinrc`` file
-        - ``PATH`` variable.
-        """
-        binary_name = self.BINARY
-
-        # Try CWD
-        binary_path = Path.cwd() / binary_name
-        if binary_path.is_file():
-            return str(binary_path)
-
-        # Try with AUSTINPATH variable
-        austin_path = os.environ.get("AUSTINPATH", None)
-        if austin_path:
-            binary_path = Path(austin_path).expanduser() / binary_name
-            if binary_path.is_file():
-                return str(binary_path)
-
-        # Try with .austinrc file
-        if AustinConfiguration().binary is not None:
-            binary_path = Path(AustinConfiguration().binary).expanduser()
-            if binary_path.is_file():
-                return str(binary_path)
-
-        # Try from PATH
-        return self.BINARY
-
-    @property
-    def version(self) -> SemVer:
-        """Austin version."""
-        return _to_semver(self._meta.get("austin"))
-
-    @property
-    def python_version(self) -> SemVer:
-        """The version of the detected Python interpreter."""
-        return _to_semver(self._meta.get("python"))

austin/aio.py

@@ -23,13 +23,15 @@
 
 import asyncio
 import sys
-from typing import Dict
-from typing import List
-from typing import Optional
+import typing as t
 
-from austin import AustinError
-from austin import BaseAustin
+from austin.base import AustinState
+from austin.base import BaseAustin
 from austin.cli import AustinArgumentParser
+from austin.errors import AustinError
+from austin.events import AustinMetadata
+from austin.events import AustinSample
+from austin.format.mojo import AsyncMojoStreamReader
 
 
 class AsyncAustin(BaseAustin):
@@ -38,69 +40,89 @@ class AsyncAustin(BaseAustin):
     Implements an ``asyncio`` API for Austin so that it can be used alongside
     other asynchronous tasks.
 
-    The following example shows how to make a simple asynchronous echo
-    implementation of Austin that behaves exactly just like Austin.
+    The following example shows how to make a simple asynchronous implementation
+    of Austin that returns stack in the collapsed format
 
     Example::
 
-        class EchoAsyncAustin(AsyncAustin):
-            def on_ready(self, process, child_process, command_line):
-                print(f"Austin PID: {process.pid}")
-                print(f"Python PID: {child_process.pid}")
-                print(f"Command Line: {command_line}")
+        import asyncio
+        import sys
 
-            def on_sample_received(self, line):
-                print(line)
+        from austin.aio import AsyncAustin
+        from austin.format.collapsed_stack import AustinEventCollapsedStackFormatter
+
+        FORMATTER = AustinEventCollapsedStackFormatter()
+
+
+        class CollapsedStackAsyncAustin(AsyncAustin):
+            async def on_sample(self, sample):
+                print(FORMATTER.format(sample))
+
+            async def on_metadata(self, metadata):
+                print(FORMATTER.format(metadata))
 
-            def on_terminate(self, data):
-                print(data)
 
         if sys.platform == "win32":
             asyncio.set_event_loop(asyncio.ProactorEventLoop())
 
-        try:
-            austin = EchoAsyncAustin()
-            asyncio.get_event_loop().run_until_complete(
-                austin.start(["-i", "10000", "python3", "myscript.py"])
-            )
-        except (KeyboardInterrupt, asyncio.CancelledError):
-            pass
-    """
 
-    async def _read_stderr(self) -> Optional[str]:
-        assert self.proc.stderr is not None
+        async def main():
+            austin = CollapsedStackAsyncAustin()
+            await austin.start(["-i", "10ms", "python", "myscript.py"])
+            await austin.wait()
 
-        try:
-            return (
-                (await asyncio.wait_for(self.proc.stderr.read(), 0.1)).decode().rstrip()
-            )
-        except asyncio.TimeoutError:
-            return None
 
-    async def _read_meta(self) -> Dict[str, str]:
-        assert self.proc.stdout is not None
-
-        meta = {}
+        asyncio.run(main())
+    """
 
-        while True:
-            line = (await self.proc.stdout.readline()).decode().rstrip()
-            if not (line and line.startswith("# ")):
-                break
-            key, _, value = line[2:].partition(": ")
-            meta[key] = value
+    def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
+        super().__init__(*args, **kwargs)
+
+        self._run_task: t.Optional[asyncio.Task] = None
+        self._proc: t.Optional[asyncio.subprocess.Process] = None
+
+    async def _run(self, mojo: AsyncMojoStreamReader) -> None:
+        # Start collecting samples
+        self._state = AustinState.RUNNING
+
+        async for e in mojo:
+            if isinstance(e, AustinSample):
+                try:
+                    await t.cast(t.Awaitable[None], self._sample_callback(e))
+                except Exception as exc:
+                    raise AustinError("Error in call to sample callback") from exc
+            elif isinstance(e, AustinMetadata):
+                self._meta[e.name] = e.value
+                if self._metadata_callback is not None:
+                    try:
+                        await t.cast(t.Awaitable[None], self._metadata_callback(e))
+                    except Exception as exc:
+                        raise AustinError("Error in call to metadata callback") from exc
+
+        self._state = AustinState.TERMINATING
+
+        # Call the terminate callback
+        try:
+            if self._terminate_callback is not None:
+                await t.cast(t.Awaitable[None], self._terminate_callback())
+        except Exception as exc:
+            raise AustinError("Error in call to terminate callback") from exc
 
-        self._meta.update(meta)
-        return meta
+        self._state = AustinState.TERMINATED
 
-    async def start(self, args: Optional[List[str]] = None) -> None:
+    async def start(self, args: t.Optional[t.Sequence[str]] = None) -> None:
         """Create the start coroutine.
 
         Use with the ``asyncio`` event loop.
         """
+        self._args = AustinArgumentParser().parse_args(args)
+
+        self._state = AustinState.STARTING
+
         try:
             _args = list(args if args is not None else sys.argv[1:])  # Make a copy
             _args.insert(0, "-P")
-            self.proc = await asyncio.create_subprocess_exec(
+            self._proc = await asyncio.create_subprocess_exec(
                 self.binary_path,
                 *_args,
                 stdin=asyncio.subprocess.PIPE,
@@ -110,44 +132,53 @@ class AsyncAustin(BaseAustin):
         except FileNotFoundError:
             raise AustinError("Austin executable not found.") from None
 
-        if not self.proc.stdout:
+        if not self._proc.stdout:
             raise AustinError("Standard output stream is unexpectedly missing")
-        if not self.proc.stderr:
+        if not self._proc.stderr:
             raise AustinError("Standard error stream is unexpectedly missing")
 
-        self._running = True
+        mojo = AsyncMojoStreamReader(self._proc.stdout)
 
-        try:
-            if not await self._read_meta():
-                raise AustinError("Austin did not start properly")
-
-            self.check_version()
+        # Retrieve the Austin version, then call the ready callback
+        async for e in mojo:
+            if isinstance(e, AustinMetadata):
+                self._meta[e.name] = e.value
 
-            self._ready_callback(
-                *self._get_process_info(
-                    AustinArgumentParser().parse_args(args), self.proc.pid
-                )
-            )
+                try:
+                    if self._metadata_callback is not None:
+                        await t.cast(t.Awaitable[None], self._metadata_callback(e))
+                except Exception as exc:
+                    raise AustinError("Error in call to metadata callback") from exc
 
-            # Start readline loop
-            while self._running:
-                data = (await self.proc.stdout.readline()).rstrip()
-                if not data:
+                if e.name == "austin":
+                    self._check_version()
                     break
+        else:
+            raise AustinError("Cannot determine Austin version from output")
 
-                self.submit_sample(data)
+        # Start the run task where we collect the samples
+        self._run_task = asyncio.create_task(self._run(mojo))
+
+    def terminate(self) -> None:
+        """Terminate Austin.
+
+        Stop the underlying Austin process by sending a termination signal.
+        """
+        if self._proc is None:
+            raise AustinError("Austin is not running")
+
+        try:
+            self._proc.terminate()
+        except ProcessLookupError as exc:
+            raise AustinError("Austin process already terminated") from exc
 
-            self._terminate_callback(await self._read_meta())
-            self.check_exit(await self.proc.wait(), await self._read_stderr())
+    async def wait(self) -> int:
+        if self._proc is None or self._run_task is None:
+            raise AustinError("Austin process is not running")
 
-        except Exception:
-            try:
-                self.proc.terminate()
-                await self.proc.wait()
-            except Exception:
-                # best effort
-                pass
-            raise
+        # Await the run task. This will terminate naturally when the Austin
+        # process has terminated.
+        await self._run_task
 
-        finally:
-            self._running = False
+        # Wait for the Austin process to terminate and return the exit code
+        return await self._proc.wait()