1. import ast
  2. import inspect
  3. import importlib
  4. import os
  5. import sys
  6. import traceback
  7. import tempfile
  8. import shutil
  11. def instrument_file(filepath, instrumentation_function):
  12. """
  13. Instruments a Python file with a given function.
  15. Args:
  16. filepath (str): The path to the Python file.
  17. instrumentation_function (callable): The function to instrument with.
  18. It should take the module object as an argument.
  20. Returns:
  21. bool: True if instrumentation was successful, False otherwise.
  22. """
  23. try:
  24. with open(filepath, 'r') as f:
  25. source_code = f.read()
  27. tree = ast.parse(source_code)
  29. for node in ast.walk(tree):
  30. if isinstance(node, ast.Module):
  31. instrumentation_function(module=ast.Module(body=[node.body]))
  32. return True
  34. return False # No module found
  35. except (SyntaxError, FileNotFoundError, Exception) as e:
  36. print(f"Error instrumenting {filepath}: {e}")
  37. return False
  40. def wrap_function(function):
  41. """Wraps a function to log its execution and handle exceptions gracefully."""
  43. def wrapper(*args, **kwargs):
  44. try:
  45. return function(*args, **kwargs)
  46. except Exception as e:
  47. print(f"Function {function.__name__} raised an exception: {e}")
  48. traceback.print_exc()
  49. return None # Or some default value, depending on the use case
  51. return wrapper
  54. def instrument_code(filepath, instrumentation_function):
  55. """
  56. Instruments the code in a file, handling potential errors gracefully.
  58. Args:
  59. filepath (str): The path to the Python file.
  60. instrumentation_function (callable): A function that takes a module
  61. object as an argument, which is
  62. used to modify the code.
  64. Returns:
  65. bool: True if instrumentation was successful, False otherwise.
  66. """
  67. try:
  68. # Create a temporary file to write the instrumented code
  69. with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file:
  70. with open(filepath, 'r') as source_file:
  71. shutil.copyfileobj(source_file, temp_file)
  72. temp_filepath = temp_file.name
  74. if instrument_file(temp_filepath, instrumentation_function):
  75. # Replace the original file with the instrumented one
  76. shutil.move(temp_filepath, filepath)
  77. return True
  78. else:
  79. # If instrumentation failed, clean up the temporary file
  80. os.remove(temp_filepath)
  81. return False
  82. except Exception as e:
  83. print(f"Error during instrumentation: {e}")
  84. return False
  87. def register_instrumentation(module):
  88. """Registers the instrumentation function with the module."""
  89. module.instrumentation_function = wrap_function # Assign the wrapper function
  92. if __name__ == '__main__':
  93. # Example Usage:
  94. # Define your instrumentation function
  95. def my_instrumentation(module):
  96. """
  97. Example instrumentation function.
  98. This function is called to instrument the module.
  99. """
  100. for name, obj in inspect.getmembers(module):
  101. if inspect.isfunction(obj):
  102. setattr(obj, 'instrumented', my_instrumentation) # Add an attribute to the function
  105. # Specify the file to instrument
  106. file_to_instrument = 'example.py'
  108. # Create a dummy example.py file for testing
  109. with open(file_to_instrument, 'w') as f:
  110. f.write("""
  111. def my_function(x):
  112. return x + 1
  114. def another_function(y):
  115. result = y / 0 # This will cause a ZeroDivisionError
  116. return result
  117. """)
  119. # Instrument the file
  120. if instrument_code(file_to_instrument, my_instrumentation):
  121. print(f"Successfully instrumented {file_to_instrument}")
  123. # Example of running the instrumented code
  124. import example
  125. print(example.my_function(5))
  126. try:
  127. print(example.another_function(10))
  128. except Exception as e:
  129. print(f"Caught exception: {e}")
  130. else:
  131. print(f"Failed to instrument {file_

Add your comment