Mock multiple calls of a class function - python-3.x

I have an issue with mocking multiple calls of a class function. My goal is to raise an Exception on a second function call. What's happening is that mocked function returns just a mocked object, without calling it. So in a result on first function call it returns just a function.
First example:
with mock.patch.object(FunctionClass, "function_to_mock") as mocked_function:
mocked_function.side_effect = [mock.DEFAULT, Exception()]
Second example:
original_function = FunctionClass.function_to_mock # to store reference to the original function
with mock.patch.object(FunctionClass, "function_to_mock") as mocked_function:
mocked_function.side_effect = [original_function, Exception()]

So I ended up with writing:
def _is_exception(obj):
return (
isinstance(obj, BaseException) or
isinstance(obj, type) and issubclass(obj, BaseException)
)
class MultipleSideEffectsMock:
def __init__(self, side_effects_collection):
self.side_effects_collection = side_effects_collection
def side_effects(self, *args, **kwargs):
side_effect = self.side_effects_collection.pop()
if callable(side_effect):
return side_effect(*args, **kwargs)
elif _is_exception(side_effect):
raise side_effect
return side_effect
which can be used as:
original_function = FunctionClass.function_to_mock # to store reference to the original function
with mock.patch.object(FunctionClass, "function_to_mock") as mocked_function:
mocked_function.side_effect = MultipleSideEffectsMock(side_effects_collection=[Exception(), original_function]).side_effects
I'm pretty sure that it could be archived by some cleaner solution or by pure mocks usage, but I didn't have time to investigate it deeper.

Related

Keep state in decorator

I am trying to write a #assert_logged_in decorator.
On paper, it's easy:
def assert_logged_in(meth: Callable):
logged_in = False
def wrapper(self, *args, **kwargs):
if not logged_in:
try:
test_login()
except as err: # bare except only for this example
raise NotLoggedIn from err
logged_in = True # <= that is what I do not understand ‽‽‽
return meth(self, *args, **kwargs)
return wrapper
# somewhere else
#assert_logged_in
def do_something(self):
pass
As long as I do not try to update logged_in, the code runs, but would test loggedinness every time.
If I try to update logged_in as in my example, I receive: UnboundLocalError: local variable 'logged_in' referenced before assignment. Even if I try to use global. Note that I do not want to keep the state in self (which works), because this decorator will be used across multiple classes.
I am very confused because this example:
def memoize(f):
memo = {}
def memoized_func(n):
if n not in memo:
memo[n] = f(n)
return memo[n]
return memoized_func
does work as intended and keeps state.
There is obviously something I do not understand. What could it be, and how could I get my decorator working?
From the comments, there are 2 options:
Cleanest option:
Declare logged_in normally outside the wrapper, and then again as nonlocal inside the wrapper.
Other possibility
from a wrapper, you cannot not change an external variable
but if this variable happens to be a reference (dict, collections...) then the wrapper can update its content without changing its reference.
For instance, have a dict instead of a boolean:
logged_in = {"logged_in": False}
instead of
logged_in = False
Anotherj comment five

Decorating a decorator

I tried to do some argument modification on a decorator which decorates a function.
The original code looks like
#original_decorator(arg=some_object)
def calculate(a, b):
# complex business logic
raise Exception()
where original_decorator is responsible for exception handling.
What I want to achieve is to do some temporary modification on some_object and restore it's property after function returned.
And I've tried the following
def replace_arg(arg, add_some_property):
def decorator_wrapper(decorator_func):
def decorator_inner(*decorator_args, **decorator_kwargs):
def actual_wrapper(actual_func):
def actual_inner(*actual_args, **actual_kwargs):
original = arg['func']
arg['func'] = add_some_property
decorator_kwargs['arg'] = arg
result = actual_func(*actual_args, **actual_kwargs)
arg['func'] = original
return result
return actual_inner
return actual_wrapper
return retry_inner
return retry_wrapper
Also tried to place the modification logic in decorator_inner, but neither worked.
My Questions:
Is it possible to modify a decorator's argument?
If true, then how can I achieve it?

how to use python decorator with argument?

I would like to define a decorator that will register classes by a name given as an argument of my decorator. I could read from stackoverflow and other sources many examples that show how to derive such (tricky) code but when adapted to my needs my code fails to produce the expected result. Here is the code:
import functools
READERS = {}
def register(typ):
def decorator_register(kls):
#functools.wraps(kls)
def wrapper_register(*args, **kwargs):
READERS[typ] = kls
return wrapper_register
return decorator_register
#register(".pdb")
class PDBReader:
pass
#register(".gro")
class GromacsReader:
pass
print(READERS)
This code produces an empty dictionary while I would expect a dictionary with two entries. Would you have any idea about what is wrong with my code ?
Taking arguments (via (...)) and decoration (via #) both result in calls of functions. Each "stage" of taking arguments or decoration maps to one call and thus one nested functions in the decorator definition. register is a three-stage decorator and takes as many calls to trigger its innermost code. Of these,
the first is the argument ((".pdb")),
the second is the class definition (#... class), and
the third is the class call/instantiation (PDBReader(...))
This stage is broken as it does not instantiate the class.
In order to store the class itself in the dictionary, store it at the second stage. As the instances are not to be stored, remove the third stage.
def register(typ): # first stage: file extension
"""Create a decorator to register its target for the given `typ`"""
def decorator_register(kls): # second stage: Reader class
"""Decorator to register its target `kls` for the previously given `typ`"""
READERS[typ] = kls
return kls # <<< return class to preserve it
return decorator_register
Take note that the result of a decorator replaces its target. Thus, you should generally return the target itself or an equivalent object. Since in this case the class is returned immediately, there is no need to use functools.wraps.
READERS = {}
def register(typ): # first stage: file extension
"""Create a decorator to register its target for the given `typ`"""
def decorator_register(kls): # second stage: Reader class
"""Decorator to register its target `kls` for the previously given `typ`"""
READERS[typ] = kls
return kls # <<< return class to preserve it
return decorator_register
#register(".pdb")
class PDBReader:
pass
#register(".gro")
class GromacsReader:
pass
print(READERS) # {'.pdb': <class '__main__.PDBReader'>, '.gro': <class '__main__.GromacsReader'>}
If you don't actually call the code that the decorator is "wrapping" then the "inner" function will not fire, and you will not create an entry inside of READER. However, even if you create instances of PDBReader or GromacsReader, the value inside of READER will be of the classes themselves, not an instance of them.
If you want to do the latter, you have to change wrapper_register to something like this:
def register(typ):
def decorator_register(kls):
#functools.wraps(kls)
def wrapper_register(*args, **kwargs):
READERS[typ] = kls(*args, **kwargs)
return READERS[typ]
return wrapper_register
return decorator_register
I added simple init/repr inside of the classes to visualize it better:
#register(".pdb")
class PDBReader:
def __init__(self, var):
self.var = var
def __repr__(self):
return f"PDBReader({self.var})"
#register(".gro")
class GromacsReader:
def __init__(self, var):
self.var = var
def __repr__(self):
return f"GromacsReader({self.var})"
And then we initialize some objects:
x = PDBReader("Inside of PDB")
z = GromacsReader("Inside of Gromacs")
print(x) # Output: PDBReader(Inside of PDB)
print(z) # Output: GromacsReader(Inside of Gromacs)
print(READERS) # Output: {'.pdb': PDBReader(Inside of PDB), '.gro': GromacsReader(Inside of Gromacs)}
If you don't want to store the initialized object in READER however, you will still need to return an initialized object, otherwise when you try to initialize the object, it will return None.
You can then simply change wrapper_register to:
def wrapper_register(*args, **kwargs):
READERS[typ] = kls
return kls(*args, **kwargs)

Using singledispatch with custom class(CPython 3.8.2)

Let's say I want to set functions for each classes in module Named 'MacroMethods'. So I've set up singledispatch after seeing it in 'Fluent Python' like this:
#singledispatch
def addMethod(self, obj):
print(f'Wrong Object {str(obj)} supplied.')
return obj
...
#addMethod.register(MacroMethods.Wait)
def _(self, obj):
print('adding object wait')
obj.delay = self.waitSpin.value
obj.onFail = None
obj.onSuccess = None
return obj
Desired behavior is - when instance of class 'MacroMethods.Wait' is given as argument, singledispatch runs registered function with that class type.
Instead, it runs default function rather than registered one.
>>> Wrong Object <MacroMethods.Wait object at 0x0936D1A8> supplied.
However, type() clearly shows instance is class 'MacroMethods.Wait', and dict_keys property also contains it.
>>> dict_keys([<class 'object'>, ..., <class 'MacroMethods.Wait'>])
I suspect all custom classes I made count as 'object' type and don't run desired functions in result.
Any way to solve this problem? Entire codes are here.
Update
I've managed to mimic singledispatch's actions as following:
from functools import wraps
def state_deco(func_main):
"""
Decorator that mimics singledispatch for ease of interaction expansions.
"""
# assuming no args are needed for interaction functions.
func_main.dispatch_list = {} # collect decorated functions
#wraps(func_main)
def wrapper(target):
# dispatch target to destination interaction function.
nonlocal func_main
try:
# find and run callable for target
return func_main.dispatch_list[type(target)]()
except KeyError:
# If no matching case found, main decorated function will run instead.
func_main()
def register(target):
# A decorator that register decorated function to main decorated function.
def decorate(func_sub):
nonlocal func_main
func_main.dispatch_list[target] = func_sub
def register_wrapper(*args, **kwargs):
return func_sub(*args, **kwargs)
return register_wrapper
return decorate
wrapper.register = register
return wrapper
Used like:
#state_deco
def general():
return "A's reaction to undefined others."
#general.register(StateA)
def _():
return "A's reaction of another A"
#general.register(StateB)
def _():
return "A's reaction of B"
But still it's not singledispatch, so I find this might be inappropriate to post this as answer.
I wanted to do similar and had the same trouble. Looks like we have bumped into a python bug. Found a write-up that describes this situation.
Here is the link to the Python Bug Tracker.
Python 3.7 breaks on singledispatch_function.register(pseudo_type), which Python 3.6 accepted

Using LogRecordFactory in python to add custom fields for logging

I am trying to add a custom field in my logging using LogRecordFactory. I am repeatedly calling a class and every time I do that, I want to set the custom_attribute in the init module so the remainder of the code within the class will have this attribute. But I cannot get this to work. I found the following which works, but its static.
import logging
old_factory = logging.getLogRecordFactory()
def record_factory(*args, **kwargs):
record = old_factory(*args, **kwargs)
record.custom_attribute = "whatever"
return record
logging.basicConfig(format="%(custom_attribute)s - %(message)s")
logging.setLogRecordFactory(record_factory)
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
logging.debug("test")
This will output correctly:
whatever - test
However, my use case is that the custom_attribute will vary. Every time I call a specific function, I want to change this. So it seems like record_factory needs another parameter passed to it so it can then return the correct record with the new parameter. But I cant figure it out. I have tried adding a parameter to the function, but when I make the call I get:
TypeError: __init__() missing 7 required positional arguments: 'name', 'level', 'pathname', 'lineno', 'msg', 'args', and 'exc_info'
I think this has something to do with the *args and **kwargs but I don't really know. Also, why are there no parenthesis after record_factory when its called by logging.setLogRecordFactory? I have never seen a function work like this.
You can try to use closure:
import logging
old_factory = logging.getLogRecordFactory()
def record_factory_factory(context_id):
def record_factory(*args, **kwargs):
record = old_factory(*args, **kwargs)
record.custom_attribute = context_id
return record
return record_factory
logging.basicConfig(format="%(custom_attribute)s - %(message)s")
logging.setLogRecordFactory(record_factory_factory("whatever"))
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
logging.debug("test")
logging.setLogRecordFactory(record_factory_factory("whatever2"))
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
logging.debug("test")
result:
$ python3 log_test.py
whatever - test
whatever2 - test
I stumbled upon this question while I was trying to do something similar. This is how I solved it, assuming that you want to add something called xyz to every log line (further explanation below):
import logging
import threading
thread_local = threading.local()
def add_xyz_to_logrecords(xyz):
factory = logging.getLogRecordFactory()
if isinstance(factory, XYZLogFactory):
factory.set_xyz(xyz)
else:
logging.setLogRecordFactory(XYZLogFactory(factory, xyz))
class XYZLogFactory():
def __init__(self, original_factory, xyz):
self.original_factory = original_factory
thread_local.xyz = xyz
def __call__(self, *args, **kwargs):
record = self.original_factory(*args, **kwargs)
try:
record.xyz = thread_local.xyz
except AttributeError:
pass
return record
def set_xyz(self, xyz):
thread_local.xyz = xyz
Here I've created a callable class XYZLogFactory, that remembers what the current value of xyz is, and also remembers what the original LogRecordFactory was. When called as a function, it creates a record using the original LogRecordFactory, and adds an xyz attribute with the current value.
The thread_local is to make it thread-safe, but for an easier version, you could just use an attribute on the XYZLogFactory:
class XYZLogFactory():
def __init__(self, original_factory, xyz):
self.original_factory = original_factory
self.xyz = xyz
def __call__(self, *args, **kwargs):
record = self.original_factory(*args, **kwargs)
record.xyz = self.xyz
return record
def set_xyz(self, xyz):
self.xyz = xyz
In my very first attempt (not shown here), I did not store the original factory, but stored it implicitly in the new LogRecordFactury using a closure. However, after a while that led to a RecursionError, because it kept calling the previous factory, which called the previous factory, etc.
Regarding your last question: there are no parentheses because the function is not called here. Instead it's passed to the logging.setLogRecordFactory, which saves it in a variable somewhere, and then calls that someplace else. If you want more information you can google something like 'functions as first class citizens'.
Easy example:
x = str # Assign to x the function that gives string representation of object
x(1) # outputs the string representation of 1, same as if you'd called str(1)
> '1'

Resources