add

server/.venv/lib/python3.11/site-packages/sqlalchemy/testing/profiling.py

@@ -0,0 +1,317 @@
+# testing/profiling.py
+# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+# mypy: allow-untyped-calls
+#
+# allow-untyped-calls is for sqlalchemy.testing.config, which is itself
+# not typed.
+
+
+"""Profiling support for unit and performance tests.
+
+These are special purpose profiling methods which operate
+in a more fine-grained way than nose's profiling plugin.
+
+"""
+
+from __future__ import annotations
+
+import collections
+import contextlib
+import os
+import platform
+import pstats
+import re
+import sys
+from typing import Any
+from typing import Callable
+from typing import DefaultDict
+from typing import Generator
+from typing import Optional
+from typing import Tuple
+
+from . import config
+from . import provision
+from .profiles_file import PlatformKey
+from .profiles_file import ProfilesFile
+from .profiles_file import TestKey
+from .util import gc_collect
+from ..util import freethreading
+from ..util import has_compiled_ext
+
+try:
+    import cProfile
+except ImportError:
+    cProfile = None  # type: ignore[assignment]
+
+_profile_stats: ProfileStatsFile = None  # type: ignore[assignment]
+"""global ProfileStatsFileInstance.
+
+plugin_base assigns this at the start of all tests.
+
+"""
+
+
+_current_test: TestKey = None  # type: ignore[assignment]
+"""String id of current test.
+
+plugin_base assigns this at the start of each test using
+_start_current_test.
+
+"""
+
+
+def _start_current_test(id_: TestKey) -> None:
+    global _current_test
+    _current_test = id_
+
+    if _profile_stats.force_write:
+        _profile_stats.reset_count()
+
+
+class ProfileStatsFile:
+    """Store per-platform/fn profiling results in a file.
+
+    There was no json module available when this was written, but now
+    the file format which is very deterministically line oriented is kind of
+    handy in any case for diffs and merges.
+
+    The parsing and rendering of the file itself, as well as of the platform
+    keys used within it, is in :mod:`.testing.profiles_file`, which is shared
+    with ``tools/profiles.py``.  This class adds the bookkeeping needed while
+    the test suite runs.
+
+    """
+
+    def __init__(
+        self,
+        filename: str,
+        sort: str = "cumulative",
+        dump: Optional[str] = None,
+    ) -> None:
+        self.force_write = (
+            config.options is not None and config.options.force_write_profiles
+        )
+        self.write = self.force_write or (
+            config.options is not None and config.options.write_profiles
+        )
+        self.profiles = ProfilesFile(filename)
+        self.fname = self.profiles.path
+        self.short_fname = os.path.split(self.fname)[-1]
+        self.dump = dump
+        self.sort = sort
+
+        # (test key, platform key) -> number of counts consumed so far
+        # within the current test
+        self._current_counts: DefaultDict[Tuple[TestKey, PlatformKey], int] = (
+            collections.defaultdict(int)
+        )
+
+        if self.write:
+            # rewrite for the case where features changed,
+            # etc.
+            self._write()
+
+    @property
+    def platform_key(self) -> PlatformKey:
+        db = config.db
+        assert db is not None, "no database configured"
+        return self.platform_key_for(db)
+
+    def platform_key_for(self, db: Any) -> PlatformKey:
+        dbapi_flags = []
+        if db.dialect.is_async:
+            dbapi_flags.append("async")
+
+        # anything else that moves the counts and that only the dialect
+        # knows about, e.g. whether the DBAPI loaded its C accelerator
+        dbapi_flags.extend(provision.profile_platform_tokens(db))
+
+        # major.minor only; the micro version isn't part of the key
+        py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+        if freethreading:
+            py_version += "t"
+
+        return PlatformKey(
+            machine=platform.machine(),
+            system=platform.system().lower(),
+            implementation=platform.python_implementation().lower(),
+            python=py_version,
+            dialect=db.name,
+            driver=db.driver,
+            dbapi_flags=tuple(dbapi_flags),
+            cextensions=has_compiled_ext(),
+        )
+
+    def has_platform(self, db: Any) -> bool:
+        """Return True if profiling tests should run against the given
+        engine.
+
+        This is the case when the profile file has counts for at least
+        one test on the engine's platform, or when profiles are being
+        written.  Used to keep profiling test classes from being generated
+        at all for drivers / interpreters that would only have each
+        test skip from within :func:`.count_functions`, after the class
+        has already run its fixture setup.
+
+        """
+        if self.write:
+            return True
+
+        platform_key = self.platform_key_for(db)
+        return any(
+            platform_key in self.profiles.platforms_for(test_key)
+            for test_key in self.profiles.test_keys()
+        )
+
+    def has_stats(self) -> bool:
+        return (
+            self.profiles.counts_for(_current_test, self.platform_key)
+            is not None
+        )
+
+    def result(self, callcount: int) -> Optional[Tuple[Optional[int], int]]:
+        test_key = _current_test
+        platform_key = self.platform_key
+
+        counts = self.profiles.setdefault_counts(test_key, platform_key)
+        current_count = self._current_counts[test_key, platform_key]
+
+        result: Optional[Tuple[Optional[int], int]]
+        if len(counts) <= current_count:
+            counts.append(callcount)
+            if self.write:
+                self._write()
+            result = None
+        else:
+            result = (
+                self.profiles.lineno_for(test_key, platform_key),
+                counts[current_count],
+            )
+        self._current_counts[test_key, platform_key] = current_count + 1
+        return result
+
+    def reset_count(self) -> None:
+        counts = self.profiles.counts_for(_current_test, self.platform_key)
+        if counts is not None:
+            counts[:] = []
+
+    def replace(self, callcount: int) -> None:
+        test_key = _current_test
+        platform_key = self.platform_key
+
+        counts = self.profiles.counts_for(test_key, platform_key)
+        assert counts, "replace() called before result() recorded a count"
+
+        current_count = self._current_counts[test_key, platform_key]
+        if current_count < len(counts):
+            counts[current_count - 1] = callcount
+        else:
+            counts[-1] = callcount
+        if self.write:
+            self._write()
+
+    def _write(self) -> None:
+        print("Writing profile file %s" % self.fname)
+        self.profiles.write()
+
+
+def function_call_count(
+    variance: float = 0.05, times: int = 1, warmup: int = 0
+) -> Callable[..., Any]:
+    """Assert a target for a test case's function call count.
+
+    The main purpose of this assertion is to detect changes in
+    callcounts for various functions - the actual number is not as important.
+    Callcounts are stored in a file keyed to Python version and OS platform
+    information.  This file is generated automatically for new tests,
+    and versioned so that unexpected changes in callcounts will be detected.
+
+    """
+
+    from sqlalchemy.util import decorator
+
+    @decorator
+    def wrap(fn: Callable[..., Any], *args: Any, **kw: Any) -> Any:
+        for warm in range(warmup):
+            fn(*args, **kw)
+
+        timerange = range(times)
+        with count_functions(variance=variance):
+            for time in timerange:
+                rv = fn(*args, **kw)
+            return rv
+
+    return wrap
+
+
+@contextlib.contextmanager
+def count_functions(variance: float = 0.05) -> Generator[None, None, None]:
+    if cProfile is None:
+        raise config._skip_test_exception("cProfile is not installed")
+
+    if not _profile_stats.has_stats() and not _profile_stats.write:
+        config.skip_test(
+            "No profiling stats available on this "
+            "platform for this function.  Run tests with "
+            "--write-profiles to add statistics to %s for "
+            "this platform." % _profile_stats.short_fname
+        )
+
+    gc_collect()
+
+    pr = cProfile.Profile()
+    pr.enable()
+    # began = time.time()
+    yield
+    # ended = time.time()
+    pr.disable()
+
+    # s = StringIO()
+    stats = pstats.Stats(pr, stream=sys.stdout)
+
+    # timespent = ended - began
+    callcount = stats.total_calls  # type: ignore[attr-defined]
+
+    expected = _profile_stats.result(callcount)
+
+    if expected is None:
+        expected_count = None
+    else:
+        line_no, expected_count = expected
+
+    print("Pstats calls: %d Expected %s" % (callcount, expected_count))
+    stats.sort_stats(*re.split(r"[, ]", _profile_stats.sort))
+    stats.print_stats()
+    if _profile_stats.dump:
+        base, ext = os.path.splitext(_profile_stats.dump)
+        test_name = _current_test.split(".")[-1]
+        dumpfile = "%s_%s%s" % (base, test_name, ext or ".profile")
+        stats.dump_stats(dumpfile)
+        print("Dumped stats to file %s" % dumpfile)
+    # stats.print_callers()
+    if _profile_stats.force_write:
+        _profile_stats.replace(callcount)
+    elif expected_count:
+        deviance = int(callcount * variance)
+        failed = abs(callcount - expected_count) > deviance
+
+        if failed:
+            if _profile_stats.write:
+                _profile_stats.replace(callcount)
+            else:
+                raise AssertionError(
+                    "Adjusted function call count %s not within %s%% "
+                    "of expected %s, platform %s. Rerun with "
+                    "--write-profiles to "
+                    "regenerate this callcount."
+                    % (
+                        callcount,
+                        (variance * 100),
+                        expected_count,
+                        _profile_stats.platform_key,
+                    )
+                )