Context manager to validate data

ぐ巨炮叔叔 提交于 2019-11-29 12:30:42

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

A decorator that accepts parameters could do it pretty nicely (http://www.artima.com/weblogs/viewpost.jsp?thread=240845):

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

def validator(test):
    def wrap(func):
        def wrapped(*args, **kwargs):
            result = func(*args, **kwargs)
            if test(result):
                return result
            return None
        return wrapped
    return wrap

The regular way to use it:

@validator(max10)
def valid_input(prompt="Please enter a name: "):
    return raw_input(prompt)

print valid_input()

Applying the decorator manually looks closer to what you're asking for:

valid_input = validator(max10)(raw_input)
print valid_input("Please enter a name: ")
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!