When you write a lot of Python that is intended to mimic arithmetic that was taken from C or assembly code, one giant advantage of Python sometimes becomes a burden: The integers are arbitrary precision, and the "built-in" wrapping of fixed-width integer types is something that has to be added manually, sometimes a truly arduous process. I recently had a thought about how this can be done in a slightly less agonizing way. This is a prototype and might require some more extension in the future, but I am pretty happy with it already. The idea is to implement a decorator that will use Python's `ast` module to add a bitmask to each arithmetic operation that occurs in the code of a decorated function. Here's the code:
```python
import ast
import functools
import inspect
def masked(mask: int):
'''
Convert arithmetic operations that occur within the decorated function body in such a way that
the result is reduced using the given bitmask. All additions, subtractions, multiplications,
left shifts, and taking powers is augmented by introducing a bitwise and with the given mask.
'''
def decorator(function):
code = inspect.getsource(function)
tree = ast.parse(code, mode='exec')
class Postprocessor(ast.NodeTransformer):
name = None
def visit_BinOp(self, node: ast.BinOp):
node = self.generic_visit(node)
if not isinstance(node.op, (ast.Add, ast.Mult, ast.Sub, ast.LShift, ast.Pow)):
return node
return ast.BinOp(node, ast.BitAnd(), ast.Constant(mask))
def visit_FunctionDef(self, node: ast.FunctionDef):
node = self.generic_visit(node)
if self.name is None:
node.name = self.name = F'__wrapped_{node.name}'
for k in range(len(node.decorator_list)):
if node.decorator_list[k].func.id == masked.__name__:
del node.decorator_list[:k + 1]
break
return node
pp = Postprocessor()
fixed = ast.fix_missing_locations(pp.visit(tree))
eval(compile(fixed, function.__code__.co_filename, 'exec'))
return functools.wraps(function)(eval(pp.name))
return decorator
```
With this decorator, you can now write:
```
@masked(0xFFFF)
def test(x: int) -> int:
return x * 0xBAAD
```
This will leave all code exactly as it is except for the binary operations involving addition, multiplication, subtraction, left shift, and computing powers - all of these are converted to having one additional bitwise and operation on top.