python - Context manager to validate data

Question

I'm trying to mull over a good solution to this and nothing is coming to mind. As an exercise, I'm trying to create a context manager that will handle data validation, something like:

validation = lambda x: len(x) <= 10

with validator(validation):
    some_data = input("Please enter a name of 10 characters or less: ")

print(some_data)

# OUTPUT
>> Please enter a name of 10 characters or less: FooBarSpamEggs
>> Please enter a name of 10 characters of less: Adam
Adam

Originally I thought about doing this with unittest.mock.patch but I realized I can't call the original function while it's patched, e.g.:

def patched(validation, *args):
    while True:
        p = __builtins__.input(args) # Doesn't work
        if validation(p):
            break
    return p

with unittest.mock.patch('builtins.input', patched):
    input("Some prompt here: ")
# fails on recursion error as patched calls itself

Then I considered writing a decorator to validate a single line, but that's really only useful if you could do something like:

@validate(lambda x: int(x) == 6)
p = input("How many sides does a d6 have? ")
# can't decorate a function call

I'm hung up on this context manager idea, however. Unfortunately I don't know if a context manager has any access to its contents, or if it's limited only to its arguments. Any thoughts?

As an aside, I know I could render this functionality in a function e.g.:

def validate_input(prompt, validation, msg_if_fail=None):
    while True:
        p = input(prompt)
        if validation(p):
            break
        if msg_if_fail is not None:
            print(msg_if_fail)
    return p

But it's not as pretty. This is, as I said, an exercise more than a practical problem.

Answer 1

You could use a regular context manager to wrap unittest.mock.patch, and save a reference the original input function prior to patching it. Then you can pass the original input to your patched function:

import unittest.mock
import contextlib
from functools import partial

def patched(validation, orig_input, *args):
    while True:
        p = orig_input(*args)
        if validation(p):
            break
    return p

@contextlib.contextmanager
def validator(validate_func):
    func = partial(patched, validate_func, input)  # original input reference saved here
    patch = unittest.mock.patch('builtins.input', func)
    patch.start()
    try:
        yield 
    finally:
        patch.stop()

validation = lambda x: len(x) <= 10

Then you can use the contextmanager like this:

with validator(validation):
    x = input("10 or less: ")

x = input("10 or less (unpatched): ")
print("done")

Sample output:

10 or less: abcdefghijklmnop
10 or less: abcdefgdfgdgd
10 or less: abcdef
10 or less (unpatched): abcdefghijklmnop
done