Method Overloading Using Decorators

#programming #python

While looking into Python design patterns, I encountered a suggestion about using Python’s introspection and decorators to implement C++-style signature-based method overloading and I thought I’d try writing a simple version to see how it might work.

import inspect
from functools import wraps

method_bank = {}

def overload(inner):
  name = inner.__qualname__
  classname = name.split(".")
  sig = [p.annotation for p in inspect.signature(inner).parameters.values()]
  if name not in method_bank:
    method_bank[name] = []
  method_bank[name].append({
    "arity": len(sig),
    "sig": sig,
    "fn": inner,
  })

  # TODO: How would you get kwargs to work here?
  @wraps(inner)
  def wrapper(*args):
    fn = None
    arity = len(args)
    for method in method_bank[name]:
      if fn:
        break
      # Not the right number of arguments
      if arity != method["arity"]:
        continue
      fn = method["fn"]
      for i in range(arity):
        argtype = method["sig"][i]
        arg = args[i]
        if argtype != inspect._empty and not isinstance(arg, argtype):
          fn = None
    if not fn:
      raise Exception(f"Invalid arguments for overloaded function {name}")
    return fn(*args)

  return wrapper

class Overloaded:
  @overload
  def fn(self, name: str):
    print(f"You supplied a string of {name}")

  @overload
  def fn(self, x: int):
    print(f"You supplied an integer of {x}")

  @overload
  def fn(self, name: str, x: int):
    print(f"You supplied both a string of {name} and an integer of {x}")

overloaded = Overloaded()

overloaded.fn("what")
overloaded.fn(1)
overloaded.fn("what", 1)
overloaded.fn(1, 2, 3)