Merge pull request #5 from P403n1x87/devel

feat!: add support for Austin 3

.flake8

@@ -1,9 +1,10 @@
 [flake8]
 select = ANN,B,B9,C,D,E,F,W,I
-ignore = ANN101,D100,D104,D107,E203,E501,W503,W606
+ignore = ANN101,ANN102,B950,D100,D104,D107,E203,E501,I001,I005,W503,W606
 max-line-length = 80
 docstring-convention = google
 import-order-style = google
+application-import-names = austin
 per-file-ignores =
   test/*:ANN,D
   noxfile.py:ANN,D

.travis.yml

@@ -25,7 +25,7 @@ jobs:
 before_script:
   # Install required Python versions
   - sudo add-apt-repository ppa:deadsnakes/ppa -y
-  - sudo apt-get -y install python3.{6..9} python3.{6..9}-dev python3.9-venv
+  - sudo apt-get -y install python3.{6..10} python3.{6..10}-dev python3.{9..10}-venv
 
   # Clone Austin development branch
   - git clone --branch devel --depth 1 https://github.com/P403n1x87/austin.git ../austin

austin/__init__.py

@@ -22,21 +22,43 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 
-from abc import ABC, abstractmethod
+from abc import ABC
+from abc import abstractmethod
 import argparse
 import functools
+from itertools import takewhile
 import os
 import os.path
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
-from austin.config import AustinConfiguration
 import psutil
 
+from austin.config import AustinConfiguration
 
 try:
-    _cached = functools.cache
+    _cached_property = functools.cached_property  # type: ignore[attr-defined]
 except AttributeError:
-    _cached = functools.lru_cache(maxsize=1)
+
+    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)
+
+    return (  # type: ignore[return-value]
+        tuple(
+            int(_)
+            for _ in "".join(
+                list(takewhile(lambda _: _.isdigit() or _ == ".", version))
+            ).split(".")
+        )
+        + (0, 0, 0)
+    )[:3]
 
 
 class AustinError(Exception):
@@ -70,12 +92,13 @@ class BaseAustin(ABC):
     """
 
     BINARY = "austin"
+    BINARY_VERSION = (3, 0, 0)
 
     def __init__(
         self,
         sample_callback: Callable[[str], None] = None,
         ready_callback: Callable[[psutil.Process, psutil.Process, str], None] = None,
-        terminate_callback: Callable[[str], None] = None,
+        terminate_callback: Callable[[Dict[str, str]], None] = None,
     ) -> None:
         """The ``BaseAustin`` constructor.
 
@@ -86,9 +109,10 @@ def __init__(
         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
+            self._sample_callback = sample_callback or self.on_sample_received  # type: ignore[attr-defined]
         except AttributeError:
             raise AustinError("No sample callback given or implemented.")
 
@@ -199,9 +223,26 @@ def submit_sample(self, data: bytes) -> None:
 
         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):
+                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: str) -> Any:
+    def on_terminate(self, stats: Dict[str, str]) -> Any:
         """Terminate event callback.
 
         Implement to be notified when Austin has terminated gracefully. The
@@ -228,8 +269,7 @@ def on_ready(
 
     # ---- Properties ----
 
-    @property
-    @_cached
+    @_cached_property
     def binary_path(self) -> str:
         """Discover the path of the Austin binary.
 
@@ -263,11 +303,11 @@ def binary_path(self) -> str:
         return self.BINARY
 
     @property
-    def version(self) -> Optional[str]:
+    def version(self) -> SemVer:
         """Austin version."""
-        return self._version
+        return _to_semver(self._meta.get("austin"))
 
     @property
-    def python_version(self) -> Optional[str]:
+    def python_version(self) -> SemVer:
         """The version of the detected Python interpreter."""
-        return self._python_version
+        return _to_semver(self._meta.get("python"))

austin/aio.py

@@ -23,9 +23,10 @@
 
 import asyncio
 import sys
-from typing import List
+from typing import Dict, List, Optional
 
-from austin import AustinError, AustinTerminated, BaseAustin
+from austin import AustinError
+from austin import BaseAustin
 from austin.cli import AustinArgumentParser
 
 
@@ -64,26 +65,42 @@ def on_terminate(self, data):
             pass
     """
 
-    async def _read_header(self) -> bool:
-        while self._python_version is None:
-            line = (await self.proc.stderr.readline()).decode().rstrip()
-            if not line:
-                return False
-            if " austin version: " in line:
-                _, _, self._version = line.partition(": ")
-            elif " Python version: " in line:
-                _, _, self._python_version = line.partition(": ")
-        return True
+    async def _read_stderr(self) -> Optional[str]:
+        assert self.proc.stderr is not None
+
+        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 = {}
+
+        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
+
+        self._meta.update(meta)
+        return meta
 
     async def start(self, args: List[str] = None) -> None:
         """Create the start coroutine.
 
         Use with the ``asyncio`` event loop.
         """
         try:
+            _args = list(args or sys.argv[1:])
+            _args.insert(0, "-P")
             self.proc = await asyncio.create_subprocess_exec(
                 self.binary_path,
-                *(args or sys.argv[1:]),
+                *_args,
                 stdin=asyncio.subprocess.PIPE,
                 stdout=asyncio.subprocess.PIPE,
                 stderr=asyncio.subprocess.PIPE,
@@ -99,10 +116,11 @@ async def start(self, args: List[str] = None) -> None:
         self._running = True
 
         try:
-            if not await self._read_header():
+            if not await self._read_meta():
                 raise AustinError("Austin did not start properly")
 
-            # Austin started correctly
+            self.check_version()
+
             self._ready_callback(
                 *self._get_process_info(
                     AustinArgumentParser().parse_args(args), self.proc.pid
@@ -111,28 +129,19 @@ async def start(self, args: List[str] = None) -> None:
 
             # Start readline loop
             while self._running:
-                data = await self.proc.stdout.readline()
+                data = (await self.proc.stdout.readline()).rstrip()
                 if not data:
                     break
 
                 self.submit_sample(data)
 
+            self._terminate_callback(await self._read_meta())
+            self.check_exit(await self.proc.wait(), await self._read_stderr())
+
+        except Exception:
+            self.proc.terminate()
+            await self.proc.wait()
+            raise
+
         finally:
-            # Wait for the subprocess to terminate
             self._running = False
-
-            try:
-                stderr = (
-                    (await asyncio.wait_for(self.proc.stderr.read(), 0.1))
-                    .decode()
-                    .rstrip()
-                )
-            except asyncio.TimeoutError:
-                stderr = ""
-            self._terminate_callback(stderr)
-
-            rcode = await self.proc.wait()
-            if rcode:
-                if rcode in (-15, 15):
-                    raise AustinTerminated(stderr)
-                raise AustinError(f"({rcode}) {stderr}")

austin/cli.py

@@ -21,7 +21,9 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from argparse import ArgumentParser, Namespace, REMAINDER
+from argparse import ArgumentParser
+from argparse import Namespace
+from argparse import REMAINDER
 from typing import Any, Callable, List, NoReturn
 
 from austin import AustinError
@@ -67,7 +69,7 @@ def __init__(
 
         def time(units: str) -> Callable[[str], int]:
             """Parse time argument with units."""
-            base = {"us": 1, "ms": 1e3, "s": 1e6}[units]
+            base = int({"us": 1, "ms": 1e3, "s": 1e6}[units])
 
             def parser(arg: str) -> int:
                 if arg.endswith("us"):
@@ -171,7 +173,7 @@ def parser(arg: str) -> int:
                 "its arguments.",
             )
 
-    def parse_args(
+    def parse_args(  # type: ignore[override]
         self, args: List[str] = None, namespace: Namespace = None
     ) -> Namespace:
         """Parse the list of arguments.

austin/config.py

@@ -23,6 +23,7 @@
 
 
 import os.path
+from typing import Any, Dict
 
 import toml
 
@@ -38,7 +39,7 @@ class AustinConfiguration:
 
     RC = os.path.join(os.path.expanduser("~"), ".austinrc")
 
-    __borg__ = {}
+    __borg__: Dict[str, Any] = {}
 
     def __init__(self) -> None:
         self.__dict__ = self.__borg__

austin/format/__init__.py

@@ -0,0 +1,20 @@
+from enum import Enum
+
+
+class Mode(Enum):
+    """Austin profiling mode."""
+
+    CPU = 0
+    WALL = 1
+    MEMORY = 2
+    FULL = 3
+
+    @classmethod
+    def from_metadata(cls, mode: str) -> "Mode":
+        """Get mode from metadata information."""
+        return {
+            "cpu": Mode.CPU,
+            "wall": Mode.WALL,
+            "memory": Mode.MEMORY,
+            "full": Mode.FULL,
+        }[mode]

austin/format/compress.py

@@ -21,33 +21,32 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-from typing import TextIO
+from typing import Dict, TextIO
 
-from austin.stats import InvalidSample, Metrics, ZERO
+from austin.stats import AustinFileReader
 
 
-def compress(source: TextIO, dest: TextIO, counts: bool = False) -> None:
+def compress(source: AustinFileReader, dest: TextIO, counts: bool = False) -> None:
     """Compress the source stream.
 
     Aggregates the metrics on equal collapsed stacks. If ``counts`` is ``True``
     then time samples are counted by occurrence rather than by sampling
     duration.
     """
-    stats = {}
+    stats: Dict[str, int] = {}
 
     for line in source:
-        try:
-            metrics, head = Metrics.parse(line)
-            if counts and metrics.time > 1:
-                metrics = Metrics(1, metrics.memory_alloc, metrics.memory_dealloc)
-        except InvalidSample:
-            continue
-        if metrics:
-            stats[head] = stats.get(head, ZERO) + metrics
-
-    dest.writelines(
-        [head + " " + str(metrics) + "\n" for head, metrics in stats.items()]
-    )
+        head, _, metric = line.rpartition(" ")
+        if "," in metric:
+            raise RuntimeError("Cannot compress samples with full metrics.")
+
+        value = 1 if counts else int(metric)
+        if metric:
+            stats[head] = stats.get(head, 0) + value
+
+    dest.writelines([f"# {k}: {v}\n" for k, v in source.metadata.items()])
+    dest.write("\n")
+    dest.writelines([head + " " + str(metric) + "\n" for head, metric in stats.items()])
 
 
 def main() -> None:
@@ -85,7 +84,9 @@ def main() -> None:
     args = arg_parser.parse_args()
 
     try:
-        with open(args.input, "r") as fin, open(args.output or args.input, "w") as fout:
+        with AustinFileReader(args.input) as fin, open(
+            args.output or args.input, "w"
+        ) as fout:
             compress(fin, fout, args.counts)
     except FileNotFoundError:
         print(f"No such input file: {args.input}")