Search code examples
pythonobject

Change current object into object of a "sibling" class


I have the following class:

class Node():
    def __init__(self, cfg: dict, cfg_path: str):
        pass

And the following two children classes:

class Master(Node):
    def __init__(self, cfg: dict, cfg_path: str):
        pass

And

class Client(Node):
    def __init__(self, cfg: dict, cfg_path: str):
        pass

At some point, a Client object must be able to be converted into a Master object. I can do this by making the program restart itself, but I am looking for a more elegant solution.

This is the main function (I stripped some of the details):

if __name__ == "__main__":
    cfg = _load_config(args.config)
    msg = "What would you like to do? Type 'help' for help."
    if cfg["type"].lower() == "client":
        n = Client(cfg, args.config)
        while True:
            print("You are a client. " + msg)
            cmd = input()
            if cmd == "help":
                pass
            elif cmd == "extend":
                n.extend()
            elif cmd == "enroll":
                n.enroll()
            elif cmd == "get_master_public_key":
                n.get_master_public_key()
            elif cmd == "lookup":
                uuid = input("Enter UUID: ")
                n.lookup(uuid)
            elif cmd == "listen":
                print("Listening...")
                n.listen()
            else:
                print("Unknown command.")
    else:
        n = Master(cfg, args.config)
        while True:
            print("You are a master. " + msg)
            # cmd = input()
            cmd = "listen"
            if cmd == "help":
                pass
            elif cmd == "extend":
                n.extend()
            elif cmd == "listen":
                print("Listening...")
                n.listen()
            else:
                print("Unknown command.")

As you can see, depending on a value in the config file, it will create either a Client or a Master. Now, if the client is given the extend command, it must be turned into a master. I can do this by updating the value in the config file to 'master' and restarting the program, but I'd like a more elegant solution.


Solution

  • First method

    What you want is ability to create a Node object dynamically. Here is one way to do that:

    class Node:
        def __init__(self, cfg: dict, cfg_path: str):
            pass
    
        def interact(self):
            while True:
                command = input("> ").strip()
                try:
                    action = getattr(self, command)
                    # Calls the action
                    action()
                    if command == "exit":
                        break
                except AttributeError:
                    print(f"Unknown command: {command}")
    
        def exit(self):
            print("Exit")
    
    
    class Client(Node):
        def help(self):
            print("Client help")
    
        def extend(self):
            print("Client extend")
    
        def lookup(self):
            uuid = input("Enter UUID: ")
            # Do something to uuid
    
        # Add methods for enroll, get_master_public_key, ...
    
    
    class Master(Node):
        def help(self):
            print("Master help")
    
        # Add methods for listen, extend, ...
    
    
    def create_node(cfg, path):
        """Create a node dynamically based on cfg['type']."""
        cfg_type = cfg["type"].lower()
        if cfg_type == "client":
            NodeType = Client
        elif cfg_type == "master":
            NodeType = Master
        else:
            raise ValueError(f"Incorrect type: {cfg_type}")
    
        node = NodeType(cfg, path)
        return node
    
    
    # Assume you got cfg and cfg_path elsewhere
    cfg = {
        "type": "client",
        # Some other data
    }
    cfg_path = "some path"
    
    node = create_node(cfg, cfg_path)
    node.interact()
    

    Notes

    • The function create_node() will take the same parameters needed to create the node, examine its contents to determine which object to create.
    • The interact method is the same for both types, so it should be in Node, the base class
    • You can add more methods if needed to each class

    Second Method: Use the cmd Library

    It seems your script is trying to poll the user's input and call methods within your class (extend, lookup, ...). Python comes with the cmd library which makes this easy.

    import cmd
    
    
    class Node(cmd.Cmd):
        prompt = "> "
    
        def __init__(self, cfg: dict, cfg_path: str):
            super().__init__()
    
        def do_exit(self, _):
            """Implement the `exit` command to exits the loop."""
            print("bye")
            return True
    
    
    class Client(Node):
        def do_extend(self, _):
            """Implement the `extend` command."""
            print("Client extend")
    
        def do_lookup(self, uuid):
            """Implement the `lookup` command."""
            print(f"Looking up uuid: {uuid}")
    
        # Add methods for enroll, get_master_public_key, ...
    
    
    class Master(Node):
        pass
    
    
    def create_node(cfg, path):
        """Create a node dynamically based on cfg['type']."""
        cfg_type = cfg["type"].lower()
        if cfg_type == "client":
            NodeType = Client
        elif cfg_type == "master":
            NodeType = Master
        else:
            raise ValueError(f"Incorrect type: {cfg_type}")
    
        node = NodeType(cfg, path)
        return node
    
    
    # Assume you got cfg and cfg_path elsewhere
    cfg = {
        "type": "client",
        # Some other data
    }
    cfg_path = "some path"
    
    node = create_node(cfg, cfg_path)
    node.cmdloop()
    

    Here is some interaction, the prompt in this case is "> "

    $ python using_cmd.py
    > help
    
    Documented commands (type help <topic>):
    ========================================
    exit  extend  help  lookup
    
    > extend
    Client extend
    > lookup foo-bar
    Looking up uuid: foo-bar
    > exit
    bye
    

    Notes

    • In order to handle the extend command, define a method do_extend. To handle the lookup command, define a method do_lookup, and so on
    • By default, the help command is already implemented
    • I highly recommend using this library as it handles lots of details for you, leaving you concentrating on writing what matters to your script.