codekingpro/portable-devtools

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