idds/tests/test_wrapper.py

0001 #!/usr/bin/env python
0002 #
0003 # Licensed under the Apache License, Version 2.0 (the "License");
0004 # You may not use this file except in compliance with the License.
0005 # You may obtain a copy of the License at
0006 # http://www.apache.org/licenses/LICENSE-2.0OA
0007 #
0008 # Authors:
0009 # - Wen Guan, <wen.guan@cern.ch>, 2024
0010
0011
0012 """
0013 Test workflow.
0014 """
0015
0016 import functools
0017 import time
0018
0019
0020 def slow_down(func):
0021     """Sleep 1 second before calling the function"""
0022     @functools.wraps(func)
0023     def wrapper_slow_down(*args, **kwargs):
0024         time.sleep(1)
0025         return func(*args, **kwargs)
0026     return wrapper_slow_down
0027
0028
0029 @slow_down
0030 def countdown(from_number):
0031     if from_number < 1:
0032         print("Liftoff!")
0033     else:
0034         print(from_number)
0035         countdown(from_number - 1)
0036
0037
0038 def repeat(num_times):
0039     def decorator_repeat(func):
0040         @functools.wraps(func)
0041         def wrapper_repeat(*args, **kwargs):
0042             for _ in range(num_times):
0043                 value = func(*args, **kwargs)
0044             return value
0045         return wrapper_repeat
0046     return decorator_repeat
0047
0048
0049 @repeat(num_times=4)
0050 def greet(name):
0051     print(f"Hello {name}")
0052
0053
0054 def repeat1(_func=None, *, num_times=2):
0055     print(_func)
0056
0057     def decorator_repeat(func):
0058         @functools.wraps(func)
0059         def wrapper_repeat(*args, **kwargs):
0060             for _ in range(num_times):
0061                 value = func(*args, **kwargs)
0062             return value
0063         return wrapper_repeat
0064
0065     if _func is None:
0066         return decorator_repeat
0067     else:
0068         return decorator_repeat(_func)
0069
0070
0071 @repeat1
0072 def greet1(name):
0073     print(f"Hello {name}")
0074
0075
0076 @repeat1(num_times=4)
0077 def greet2(name):
0078     print(f"Hello {name}")
0079
0080
0081 if __name__ == "__main__":
0082     # countdown(10)
0083     greet1('kitty')
0084     greet2('kitty2')
0085
0086     greet2('kitty3')