codekingpro/portable-devtools

1# testing/profiling.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7# mypy: ignore-errors
8
9
10"""Profiling support for unit and performance tests.
11
12These are special purpose profiling methods which operate
13in a more fine-grained way than nose's profiling plugin.
14
15"""
16
17from __future__ import annotations
18
19import collections
20import contextlib
21import os
22import platform
23import pstats
24import re
25import sys
26
27from . import config
28from .util import gc_collect
29from ..util import freethreading
30from ..util import has_compiled_ext
31
32
33try:
34    import cProfile
35except ImportError:
36    cProfile = None
37
38_profile_stats = None
39"""global ProfileStatsFileInstance.
40
41plugin_base assigns this at the start of all tests.
42
43"""
44
45
46_current_test = None
47"""String id of current test.
48
49plugin_base assigns this at the start of each test using
50_start_current_test.
51
52"""
53
54
55def _start_current_test(id_):
56    global _current_test
57    _current_test = id_
58
59    if _profile_stats.force_write:
60        _profile_stats.reset_count()
61
62
63class ProfileStatsFile:
64    """Store per-platform/fn profiling results in a file.
65
66    There was no json module available when this was written, but now
67    the file format which is very deterministically line oriented is kind of
68    handy in any case for diffs and merges.
69
70    """
71
72    def __init__(self, filename, sort="cumulative", dump=None):
73        self.force_write = (
74            config.options is not None and config.options.force_write_profiles
75        )
76        self.write = self.force_write or (
77            config.options is not None and config.options.write_profiles
78        )
79        self.fname = os.path.abspath(filename)
80        self.short_fname = os.path.split(self.fname)[-1]
81        self.data = collections.defaultdict(
82            lambda: collections.defaultdict(dict)
83        )
84        self.dump = dump
85        self.sort = sort
86        self._read()
87        if self.write:
88            # rewrite for the case where features changed,
89            # etc.
90            self._write()
91
92    @property
93    def platform_key(self):
94        dbapi_key = config.db.name + "_" + config.db.driver
95        if config.db.dialect.is_async:
96            dbapi_key += "_async"
97
98        if config.db.name == "sqlite" and config.db.dialect._is_url_file_db(
99            config.db.url
100        ):
101            dbapi_key += "_file"
102
103        # keep it at 2.7, 3.1, 3.2, etc. for now.
104        py_version = ".".join([str(v) for v in sys.version_info[0:2]])
105        if freethreading:
106            py_version += "t"
107
108        platform_tokens = [
109            platform.machine(),
110            platform.system().lower(),
111            platform.python_implementation().lower(),
112            py_version,
113            dbapi_key,
114        ]
115
116        platform_tokens.append("dbapiunicode")
117        _has_cext = has_compiled_ext()
118        platform_tokens.append(_has_cext and "cextensions" or "nocextensions")
119        return "_".join(platform_tokens)
120
121    def has_stats(self):
122        test_key = _current_test
123        return (
124            test_key in self.data and self.platform_key in self.data[test_key]
125        )
126
127    def result(self, callcount):
128        test_key = _current_test
129        per_fn = self.data[test_key]
130        per_platform = per_fn[self.platform_key]
131
132        if "counts" not in per_platform:
133            per_platform["counts"] = counts = []
134        else:
135            counts = per_platform["counts"]
136
137        if "current_count" not in per_platform:
138            per_platform["current_count"] = current_count = 0
139        else:
140            current_count = per_platform["current_count"]
141
142        has_count = len(counts) > current_count
143
144        if not has_count:
145            counts.append(callcount)
146            if self.write:
147                self._write()
148            result = None
149        else:
150            result = per_platform["lineno"], counts[current_count]
151        per_platform["current_count"] += 1
152        return result
153
154    def reset_count(self):
155        test_key = _current_test
156        # since self.data is a defaultdict, don't access a key
157        # if we don't know it's there first.
158        if test_key not in self.data:
159            return
160        per_fn = self.data[test_key]
161        if self.platform_key not in per_fn:
162            return
163        per_platform = per_fn[self.platform_key]
164        if "counts" in per_platform:
165            per_platform["counts"][:] = []
166
167    def replace(self, callcount):
168        test_key = _current_test
169        per_fn = self.data[test_key]
170        per_platform = per_fn[self.platform_key]
171        counts = per_platform["counts"]
172        current_count = per_platform["current_count"]
173        if current_count < len(counts):
174            counts[current_count - 1] = callcount
175        else:
176            counts[-1] = callcount
177        if self.write:
178            self._write()
179
180    def _header(self):
181        return (
182            "# %s\n"
183            "# This file is written out on a per-environment basis.\n"
184            "# For each test in aaa_profiling, the corresponding "
185            "function and \n"
186            "# environment is located within this file.  "
187            "If it doesn't exist,\n"
188            "# the test is skipped.\n"
189            "# If a callcount does exist, it is compared "
190            "to what we received. \n"
191            "# assertions are raised if the counts do not match.\n"
192            "# \n"
193            "# To add a new callcount test, apply the function_call_count \n"
194            "# decorator and re-run the tests using the --write-profiles \n"
195            "# option - this file will be rewritten including the new count.\n"
196            "# \n"
197        ) % (self.fname)
198
199    def _read(self):
200        try:
201            profile_f = open(self.fname)
202        except OSError:
203            return
204        for lineno, line in enumerate(profile_f):
205            line = line.strip()
206            if not line or line.startswith("#"):
207                continue
208
209            test_key, platform_key, counts = line.split()
210            per_fn = self.data[test_key]
211            per_platform = per_fn[platform_key]
212            c = [int(count) for count in counts.split(",")]
213            per_platform["counts"] = c
214            per_platform["lineno"] = lineno + 1
215            per_platform["current_count"] = 0
216        profile_f.close()
217
218    def _write(self):
219        print("Writing profile file %s" % self.fname)
220        profile_f = open(self.fname, "w")
221        profile_f.write(self._header())
222        for test_key in sorted(self.data):
223            per_fn = self.data[test_key]
224            profile_f.write("\n# TEST: %s\n\n" % test_key)
225            for platform_key in sorted(per_fn):
226                per_platform = per_fn[platform_key]
227                c = ",".join(str(count) for count in per_platform["counts"])
228                profile_f.write("%s %s %s\n" % (test_key, platform_key, c))
229        profile_f.close()
230
231
232def function_call_count(variance=0.05, times=1, warmup=0):
233    """Assert a target for a test case's function call count.
234
235    The main purpose of this assertion is to detect changes in
236    callcounts for various functions - the actual number is not as important.
237    Callcounts are stored in a file keyed to Python version and OS platform
238    information.  This file is generated automatically for new tests,
239    and versioned so that unexpected changes in callcounts will be detected.
240
241    """
242
243    # use signature-rewriting decorator function so that pytest fixtures
244    # still work on py27.  In Py3, update_wrapper() alone is good enough,
245    # likely due to the introduction of __signature__.
246
247    from sqlalchemy.util import decorator
248
249    @decorator
250    def wrap(fn, *args, **kw):
251        for warm in range(warmup):
252            fn(*args, **kw)
253
254        timerange = range(times)
255        with count_functions(variance=variance):
256            for time in timerange:
257                rv = fn(*args, **kw)
258            return rv
259
260    return wrap
261
262
263@contextlib.contextmanager
264def count_functions(variance=0.05):
265    if cProfile is None:
266        raise config._skip_test_exception("cProfile is not installed")
267
268    if not _profile_stats.has_stats() and not _profile_stats.write:
269        config.skip_test(
270            "No profiling stats available on this "
271            "platform for this function.  Run tests with "
272            "--write-profiles to add statistics to %s for "
273            "this platform." % _profile_stats.short_fname
274        )
275
276    gc_collect()
277
278    pr = cProfile.Profile()
279    pr.enable()
280    # began = time.time()
281    yield
282    # ended = time.time()
283    pr.disable()
284
285    # s = StringIO()
286    stats = pstats.Stats(pr, stream=sys.stdout)
287
288    # timespent = ended - began
289    callcount = stats.total_calls
290
291    expected = _profile_stats.result(callcount)
292
293    if expected is None:
294        expected_count = None
295    else:
296        line_no, expected_count = expected
297
298    print("Pstats calls: %d Expected %s" % (callcount, expected_count))
299    stats.sort_stats(*re.split(r"[, ]", _profile_stats.sort))
300    stats.print_stats()
301    if _profile_stats.dump:
302        base, ext = os.path.splitext(_profile_stats.dump)
303        test_name = _current_test.split(".")[-1]
304        dumpfile = "%s_%s%s" % (base, test_name, ext or ".profile")
305        stats.dump_stats(dumpfile)
306        print("Dumped stats to file %s" % dumpfile)
307    # stats.print_callers()
308    if _profile_stats.force_write:
309        _profile_stats.replace(callcount)
310    elif expected_count:
311        deviance = int(callcount * variance)
312        failed = abs(callcount - expected_count) > deviance
313
314        if failed:
315            if _profile_stats.write:
316                _profile_stats.replace(callcount)
317            else:
318                raise AssertionError(
319                    "Adjusted function call count %s not within %s%% "
320                    "of expected %s, platform %s. Rerun with "
321                    "--write-profiles to "
322                    "regenerate this callcount."
323                    % (
324                        callcount,
325                        (variance * 100),
326                        expected_count,
327                        _profile_stats.platform_key,
328                    )
329                )
330