Daniel
Daniel

Reputation: 12026

How to only allow function to be called from the REPL, not by other functions?

Context:

I'm writing a personal python module to simplify some scripts I have lying around. One of the functions I have is untested and may have undesirable edge cases that I still have to consider. In order to not allow myself from relying on it from other modules or functions, I was wondering whether I could enforce it to raise an error if not called directly from the REPL.

I'm not asking whether this is a good idea or not. It obviously isn't because it defeats the purpose of writing a function in the first place. I'm wondering if is is possible in Python, and how to do it.

Question:

Is it possible to have a function raise an error if not called interactively? For example:

def is_called_from_top_level():
    "How to implement this?"
    pass

def shady_func():
    "Only for testing at the REPL. Calling from elsewhere will raise."
    if not is_called_from_top_level():
        raise NotImplementedError("Shady function can only be called directly.")
    return True

def other_func():
    "Has an indirect call to shady."
    return shady_func()

And then at a REPL:

[In:1] shady_func()
[Out:1] True
[In:2] other_func()
[Out:2] NotImplementedError: "Shady function can only be called directly."

Upvotes: 2

Views: 171

Answers (4)

norok2
norok2

Reputation: 26886

DISCLAIMER: This is a bit of a hack, and may not work across different Python / IPython / Jupyter versions, but the underlying idea still holds, i.e. use inspect to get an idea of who is calling.

The code below was tested with Python 3.7.3, IPython 7.6.1 and Jupyter Notebook Server 5.7.8.


Using inspect (obviously), one can look for distinctive features of the REPL frame:

  1. inside a Jupyter Notebook you can check if the repr() of the previous frame contain the string 'code <module>';
  2. using Python / IPython you can check for the code representation of the previous frame to start at line 1.

In code, this would look like:

import inspect


def is_called_from_top_level():
    "How to implement this?"
    pass


def shady_func():
    "Only for testing at the REPL. Calling from elsewhere will raise."
    frame = inspect.currentframe()
    is_interactive = (
        'code <module>' in repr(frame.f_back)  # Jupyter
        or 'line 1>' in repr(frame.f_back.f_code))  # Python / IPython
    if not is_interactive:
        raise NotImplementedError("Shady function can only be called directly.")
    return True


def other_func():
    "Has an indirect call to shady."
    return shady_func()


shady_func()
# True
other_func()
# raises NotImplementedError

(EDITED to include support for both Jupyter Notebook and Python / IPython).


As suggested by @bananafish, this is actually a good use case for a decorator:

import inspect
import functools


def repl_only(func):
    @functools.wraps(func)
    def wrapped(*args, **kwargs):
        frame = inspect.currentframe()
        is_interactive = (
            'code <module>' in repr(frame.f_back)  # Jupyter
            or 'line 1>' in repr(frame.f_back.f_code))  # Python / IPython
        if not is_interactive:
            raise NotImplementedError('Can only be called from REPL')
        return func(*args, **kwargs)
    return wrapped


@repl_only
def foo():
    return True


def bar():
    return foo()


print(foo())
# True
print(bar())
# raises NotImplementedError

Upvotes: 2

Daniel
Daniel

Reputation: 12026

Inspired by the comment to the OP suggesting looking at the stack trace, @norok2 's solution based on direct caller inspection, and by @bananafish 's use of the decorator, I came up with an alternative solution that does not require inspect nor sys. The idea is to throw and catch to get a handle on a traceback object (essentially our stack trace), and then do the direct caller inspection.

from functools import wraps


def repl_only(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        try:
            raise Exception
        except Exception as e:
            if "module" not in str(e.__traceback__.tb_frame.f_back)[-10:]:
                raise  NotImplementedError(f"{func.__name__} has to be called from the REPL!")
        return func(*args, **kwargs)
    return wrapped


@repl_only
def dangerous_util_func(a, b):
    return a + b

def foo():
    return dangerous_util_func(1, 2)

Here dangerous_util_func will run and foo will throw.

Upvotes: 0

bananafish
bananafish

Reputation: 2917

Try checking for ps1 on sys.

import sys

def dangerous_util_func(a, b):
    is_interactive = bool(getattr(sys, 'ps1', False))
    print(is_interactive)  # Prints True or False
    return a + b

You can even get fancy and make a decorator for this to make it more reusable.

import sys
from functools import wraps


def repl_only(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        is_interactive = bool(getattr(sys, 'ps1', False))
        if not is_interactive:
            raise NotImplementedError("Can only be called from REPL")
        return func(*args, **kwargs)
    return wrapped


@repl_only
def dangerous_util_func(a, b):
    return a + b

Upvotes: 4

logicOnAbstractions
logicOnAbstractions

Reputation: 2580

You can do something like that:

import inspect

def other():
    shady()

def shady():
    curfrm = inspect.currentframe()
    calframe = inspect.getouterframes(curfrm, 2)
    caller = calframe[1][3]

    if not '<module>' in caller::
        raise Exception("Not an acceptable caller")

    print("that's fine")

if __name__ == '__main__':
    import sys
    args = sys.argv[1:]
    shady()
    other()

The module inspect allows you to get information such as the function's caller. You may have to dig a bit deeper if you have edge cases....

Upvotes: 1

Related Questions