venn.py

import matplotlib.pyplot as plt
from matplotlib_venn import venn2, venn3
import upsetplot

import matplotlib_inline.backend_inline
matplotlib_inline.backend_inline.set_matplotlib_formats('svg')
%config InlineBackend.figure_format = 'retina'
%matplotlib inline

plt.rcParams['figure.figsize'] = [6, 3]
plt.rcParams['figure.dpi'] = 140
plt.rcParams['font.size'] = 8

import os, re
import traceback

def remove_special_characters(s): return re.sub('[^a-zA-Z0-9]', ' ', s)

def get_strings(stack):
    
    filename, lineno, function_name, code = stack[-2]
    names = [n.strip() for n in code[code.find("(")+1:code.rfind(")")].split(',')]

    common_prefix = os.path.commonprefix(names)
    common_suffix = os.path.commonprefix([name[::-1] for name in names])[::-1]
    if len(common_prefix) > 2:
        names = [name.removeprefix(common_prefix) for name in names]
    if len(common_suffix) > 2:
        names = [name.removesuffix(common_suffix) for name in names]
    if len(common_prefix) > 2 or len(common_suffix) > 2:
        names = [remove_special_characters(name) for name in names]    

    return names

def form_list_of_tuples(names, iterables):
    
    final = []
    for name, iterable in zip(names, iterables):

        len_full = len(iterable)
        s = set(iterable)
        len_uniques = len(s)
        if len_uniques == len_full:
            name += f' ({len_uniques})'
        else:
            name += f' ({len_uniques} unique in {len_full})'

        final.append((s, name))
    
    return final

def venn(*args, **kwargs):

    if len(args) == 2: method = venn2
    elif len(args) == 3: method = venn3
    else: print('incorrect number of args')

    names = get_strings(traceback.extract_stack())
    final = form_list_of_tuples(names, args)

    return method(*zip(*final), **kwargs)


def upset(*args, **kwargs):

    names = get_strings(traceback.extract_stack())
    final = form_list_of_tuples(names, args)

    s = upsetplot.from_contents(dict(final))
    
    fig, upsetplot.UpSet(s, subset_size='count', show_counts=True).plot()
    
    return s, fig