1

I am trying to use the monkey patching to replace a class being referenced in all the running code with the new class I have developed. I am not sure if monkey patching will create some kind of cycle here, but for now it's not working and I'm not sure why.

ssh_hook.py I want to override the get_conn method of this class, there are other methods as well. However, I just want their default implementation.

class SSHHook(object):
    def get_conn(self):
        print("SSH Hook")

custom_ssh_hook.py Class which overrides method from above class.

from airflow.test_hooks.ssh_hook import SSHHook

class CustomSSHHook(SSHHook):
    def get_conn(self):
        try:
            print("custom ssh Hook")
            return super(CustomSSHHook, self).get_conn()
        except Exception as e:
            print("retry mechanism")
            raise e

init.py

from airflow.test_hooks.ssh_hook import SSHHook

print "running ssh hook call"
ssh_hook = SSHHook()
ssh_hook.get_conn()
print "completed ssh hook call"

test_file.py this one is in airflow.test_hooks package

from airflow.test_hooks.custom_ssh_hook import CustomSSHHook
from airflow.test_hooks.ssh_hook import SSHHook

# SSHHook = CustomSSHHook
print "running monkey patching ssh hook get_conn"
SSHHook.get_conn = CustomSSHHook.get_conn
print "completed monkey patching ssh hook get_conn"

I am getting following error when I run the above code:

running monkey patching ssh hook get_conn
completed monkey patching ssh hook get_conn
running ssh hook call
Traceback (most recent call last):
  File "C:/Users/gshah03/PycharmProjects/demo/airflow/dags/test_file.py", line 5, in <module>
    ssh_hook.get_conn()
TypeError: unbound method get_conn() must be called with CustomSSHHook instance as first argument (got nothing instead)

Update:

if I monkey patch the whole class, I am getting the following output:

init.py

print "running monkey patching ssh hook get_conn"
# SSHHook.get_conn = CustomSSHHook.get_conn
SSHHook = CustomSSHHook
print "completed monkey patching ssh hook get_conn"

Output

running monkey patching ssh hook get_conn
completed monkey patching ssh hook get_conn
running ssh hook call
SSH Hook
completed ssh hook call

Update 2: following changes in __init__.py and it's working now as expected.

from airflow_lambda.test_hooks import custom_ssh_hook
from airflow_lambda.test_hooks import ssh_hook

print "running monkey patching ssh hook get_conn"
ssh_hook.SSHHook = custom_ssh_hook.CustomSSHHook
# ssh_hook.SSHHook.get_conn = custom_ssh_hook.CustomSSHHook.get_conn
print "completed monkey patching ssh hook get_conn"
Gaurang Shah
  • 11,764
  • 9
  • 74
  • 137
  • Does this answer your question? [Monkey patching a class in another module in Python](https://stackoverflow.com/questions/19545982/monkey-patching-a-class-in-another-module-in-python) – Jared Smith Feb 02 '21 at 18:15
  • 1
    Put a `print` in the `__init__.py` and `test_file.py` to make sure all the code you think is running actually is, and in the right order. Side-note: That `super` call is doubly invalid when monkey-patching `SSHHook`, because `self` is not a `CustomSSHHook` (it might work but for the next issue, but it would be by accident of implementation; you misused it) and because neither class is a new-style class (unless you omitted an assignment to `__metaclass__` or inheritance from `object` in the posted code, they're Py2's default old-style class). – ShadowRanger Feb 02 '21 at 18:16
  • @JaredSmith I looked the code. However, if I get the method out of class, I can't call the method from `SSHHook` class. I want to override that method and then replace in already running code. – Gaurang Shah Feb 02 '21 at 18:20
  • @ShadowRanger yes, object part I missed, as I was creating this sample code. :) thanks for pointing it out. – Gaurang Shah Feb 02 '21 at 18:22
  • @ShadowRanger you were right, now sure why code from `__init__.py` is not being invoked. am I missing something ? – Gaurang Shah Feb 02 '21 at 18:27

0 Answers0