# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_stata.ipynb.

# %% auto 0
__all__ = ['run_direct', 'run_single', 'get_local', 'set_local', 'get_global', 'stata_formatted', 'variable_names', 'drop_var',
           'obs_count', 'pwd', 'resolve_macro']

# %% ../nbs/02_stata.ipynb 5
from .misc_utils import DivertedPrints, print_red

# %% ../nbs/02_stata.ipynb 7
def run_direct(cmds, quietly=False, echo=False, inline=True):
    import pystata
    if quietly:
        with DivertedPrints() as diverted: # to prevent blank line output, as with `program define`
            out = pystata.stata.run(cmds, quietly, echo, inline)
            prints = diverted.getvalue()
        for line in prints.splitlines():
            if line.strip():
                print(line)
        return out
    else:
        return pystata.stata.run(cmds, quietly, echo, inline)

# %% ../nbs/02_stata.ipynb 13
def run_single(cmd, echo=False, show_exc_warning=True):
    import sfi
    try:
        sfi.SFIToolkit.stata(cmd, echo)
    except Exception as e:
        if show_exc_warning:
            print_red(f"run_single (sfi.SFIToolkit.stata) error: {repr(e)}")
        run_direct(cmd, echo=echo)

# %% ../nbs/02_stata.ipynb 27
def get_local(name):
    import sfi
    return sfi.Macro.getLocal(name)

# %% ../nbs/02_stata.ipynb 29
def set_local(name, value):
    import sfi
    return sfi.Macro.setLocal(name, value)

# %% ../nbs/02_stata.ipynb 31
def get_global(name):
    import sfi
    return sfi.Macro.getGlobal(name)

# %% ../nbs/02_stata.ipynb 33
def stata_formatted(value, s_format):
    import sfi
    return sfi.SFIToolkit.formatValue(value, s_format)

# %% ../nbs/02_stata.ipynb 35
def variable_names():
    from sfi import Data
    return [Data.getVarName(i) for i in range(Data.getVarCount())]

# %% ../nbs/02_stata.ipynb 38
def drop_var(name):
    import sfi
    sfi.Data.dropVar(name)

# %% ../nbs/02_stata.ipynb 41
def obs_count():
    """Count the number of observations"""
    import sfi
    return sfi.Data.getObsTotal()

# %% ../nbs/02_stata.ipynb 44
def pwd():
    from sfi import SFIToolkit
    return SFIToolkit.getWorkingDir()

# %% ../nbs/02_stata.ipynb 48
def resolve_macro(macro):
    macro = macro.strip()
    if macro.startswith("`") and macro.endswith("'"):
        macro = get_local(macro[1:-1])
    elif macro.startswith("$_"):
        macro = get_local(macro[2:])
    elif macro.startswith("${") and macro.endswith("}"):
        macro = get_global(macro[2:-1])
    elif macro.startswith("$"):
        macro = get_global(macro[1:])
    return macro
