get_network_driver and multiple custom_napalm modules return incorrect modules

[blargh2015]

Description of Issue/Question

When having multiple custom_napalm modules with one inheriting from another, a "get_network_driver" will return the first module alphabetically, instead of the correct module, due to this code in base/init.py get_network_driver function:

    for name, obj in inspect.getmembers(module):
        if inspect.isclass(obj) and issubclass(obj, NetworkDriver):
            return obj

inspect.getmembers returns items alphabetically, so you get the first module name imported, instead of the correct module.

Setup

napalm version

napalm==2.4.0

Network operating system version

N/A

Steps to Reproduce the Issue

Create a custom_napalm directory with an empty init.py. Add custom_napalm/a.py with:

from napalm.base.base import NetworkDriver

class ADriver(NetworkDriver):
  pass

A file custom_napalm/b.py:

from custom_napalm.a import ADriver

class BDriver(ADriver):
  print("B loaded.")

  def __init__(self, hostname, username, password, timeout=60, optional_args=None):
    print("B init called.")

And then a test.py:

import napalm
import inspect
driver = napalm.get_network_driver('b')
print(inspect.getmodule(driver))

Executing test.py:

$ python3 ./test.py
B loaded.
<module 'custom_napalm.a' from '/home/user/custom_napalm/a.py'>

(Expecting to get custom_napalm.b, not .a)

We ran into this problem because our "a.py" defines a bunch of internal functions we use (with docs, etc.) and NotImplementedErrors. Individual b.py files implement for different OSes.

[itdependsnetworks]

I do not believe this is supported. It is a convenience method provided for you within a context of a single OS per file. Any additional shared functionality would likely (I'm not a maintainer) have to be managed by yourself.

[blargh2015]

It works fine if we change the names around (i.e. change "ADriver" to "ZZZDriver"), which is what we've done internally, for the moment. Keep in mind in our real use, "a.py" doesn't define any functionality, just a bunch of functions and the docstrings for them, so we don't have to re-write that for each OS (which we have in custom_napalm/iosxr.py, custom_napalm/junos.py, etc. all of which import a.py)

[blargh2015]

We've tinkered with replacing the code in get_network_driver from:

    for name, obj in inspect.getmembers(module):
        if inspect.isclass(obj) and issubclass(obj, NetworkDriver):
            return obj

To:

for name, obj in inspect.getmembers(module, inspect.isclass):
    if obj.__module__ == module_name and issubclass(obj, NetworkDriver):
        return obj

Which works for this use case, but this may break other use cases (it can only return a class defined in the exact module loaded).

[bewing]

I think the proposed fix makes a certain kind of sense. I hadn't run into this issue as my internal implementations sub-classes NetworkDriver to OurNetworkDriver and we re-define get_network_driver to test if the other classes are subclasses of OutNetworkDriver, which may be another (especially if you want to be sure that the class you are loading has any extra methods defined in ADriver

[ktbyers]

Okay, let me try to look at it. On quick glance I suspect the quick fix isn't viable as it would probably break existing code.

This is probably something we should support in some way.