patched function.py

from types import CodeType, FunctionType
from dis import dis

def patch(fn, name, codestring):
    code = fn.__code__
    code = CodeType(
        code.co_argcount, code.co_posonlyargcount, code.co_kwonlyargcount,
        code.co_nlocals, code.co_stacksize, code.co_flags, codestring,
        code.co_consts, code.co_names, code.co_varnames,
        code.co_filename, name, code.co_firstlineno, code.co_lnotab,
        code.co_freevars,  code.co_cellvars,
    )
    return FunctionType(
        code, fn.__globals__, name, fn.__defaults__, fn.__closure__
    )

def only_dict(d, i):
    d[i] = 0
    return d[i]

def without_walrus(d, i):
    r = 0
    d[i] = r
    return r

def with_walrus(d, i):
    d[i] = (r := 0)
    return r

patched_walrus = patch(with_walrus, 'patched_walrus', bytes([100, 1, 4, 0, 124, 0, 124, 1, 60, 0, 83, 0]))

dis(patched_walrus)

plot.py

import numpy as np
import pandas as pd
from scipy import stats
import matplotlib.pyplot as plt
import seaborn as sns

sns.set(style='ticks')

df = pd.read_csv('idx_vs_walrus.csv')

fig, (ax1, ax2) = plt.subplots(
    ncols=2, sharey=True, figsize=(9, 4),
    gridspec_kw=dict(width_ratios=[10, 1]),
)

mn, mx = np.quantile(df.values, [0, 0.99])
d = mx - mn
ax1.set_ylim(mn-0.06*d, mx+0.15*d)
ax1.set_xmargin(0)

for name, series in df.items():
    sd = series.values.copy()
    s2 = series.sort_values()[95:]
    s1 = sd; s1[s2.index] = np.nan
    
    al, = ax1.plot(s1, label=name)
    ax1.plot(s2, '.', color=al.get_color())

    s1 = s1[~np.isnan(s1)]
    mu = np.mean(s1)
    sd = np.std(s1)
    std = sd / np.sqrt(sum(s1))

    print(*np.array([
        mu-2*std, mu, mu+2*std,
    ]).round(4))

    ax1.axhline(mu, ls='--', color='black', lw=1, alpha=0.5)
    kde = stats.gaussian_kde(s1)
    x = np.linspace(mu-6*sd, mu+6*sd, 201)

    ax2.plot(kde.evaluate(x), x, color=al.get_color())

ax1.set_xlabel('Timer run')
ax1.set_ylabel('time (seconds)')

ax1.legend()
ax2.get_xaxis().set_visible(False)

sns.despine(ax=ax1)
sns.despine(ax=ax2, bottom=True)

fig.tight_layout()

fig.savefig('idx_vs_walrus.png', dpi=160)

Thank you !

no probs, happy to help

Do you have any documentation that explain how to do what you do in patched walrus ? How it works, etc ?

nothing apart from the documentation that's already on line. I'd suggest looking at Python's dis module to know what it's doing. i.e. try disassembling the functions in patched function.py to see what OP codes they get compiled down to

my code in patch was made up while answering the question. I've never tried to do anything like it before so there might well be better ways to go about doing it

Thank you :)