sharing just the stupid ensure optimizations util

utils.py

#   Copyright 2023-present Michael Hall
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import ast
import inspect
import sys
from typing import TypeVar

T = TypeVar("T")

_cycle_blocked = False

def ensure_optimizations(f: T) -> T:
    """ ensures optimizations as a decorator on a pure python class or function """

    # Prevents a cyclic compile issue
    global _cycle_blocked  # noqa: PLW0603
    if _cycle_blocked:
        return f

    new_ast = compile(
        ast.parse(inspect.getsource(f)),  # type: ignore
        '<string>',
        'exec',
        optimize=1,
    )

    env = sys.modules[f.__module__].__dict__

    _cycle_blocked = True  # noqa: PLW0603
    try:
        exec(new_ast, env)  # noqa: S102
    finally:
        _cycle_blocked = False

    return env[f.__name__]  # type: ignore