MyPy and Dict-Like get() Methods

Last updated:

I'm hacking on something for my `kdl-py` project, and realized I wanted a .get() method to retrieve child nodes by name, similar to a dict. I had some trouble working out how to set it up to type properly in MyPy, so here's a short summary of my results, which I'm quite happy with now.

(Disclaimer: while working on this, I was laboring under the misapprehension that dict.get("foo") raised an exception if "foo" wasn't in the dict, but that's only the behavior for dict["foo"]! dict.get("foo") will instead always return the default value if the key is missing, which just defaults to None, which is a much simpler behavior, ugh.)

So, the problem I was having is that I wanted an optional argument (the default value, to be returned if the node name couldn't be found), and I wanted to tell whether that argument was passed at all. Any value is valid for the default, so I can't rely on a standard sentinel value, like None.

One way to do this is with kwargs shenanigans (leaving the argument out of the arglist, using a **kwargs arg instead, and just checking if it shows up in there), but that's awkward at the best of times, and doesn't let you typecheck well (can't indicate that the call might return the default value, type-wise).

The usual way to do this in JS, which doesn't have a kwargs equivalent, is instead to set the default value to some unique object value that's not exposed to the outside, and see if it's still equal to that value. Since the outside world doesn't have access to that value, you can be sure that if you see it, the argument wasn't passed at all.

This is how I ended up going. Here's the final code:

import typing as t

class _MISSING:
  pass

T = t.TypeVar('T')

class NodeList:
  @t.overload
  def get(self, key: str) -> Node:
    ...

  @t.overload
  def get(
    self, 
    key: str, 
    default: t.Union[T, _MISSING] = _MISSING(),
  ) -> t.Union[Node, T]:
    ...

  def get(
    self,
    key: str,
    default: t.Union[T, _MISSING] = _MISSING(),
  ) -> t.Union[Node, T]:
    if self.data.has(key):
      return self.data.get(key)
    if isinstance(default, _MISSING):
      raise KeyError(f"No node with name '{key}' found.")
    else:
      return default

Boom, there you go. Now if you call nl.get("foo"), the return type is definitely Node, so you don't have to do a None check to satisfy MyPy (it'll just throw if you screw up), but it'll correctly type as "Node or whatever your default is" when you do pass a default value.

(a limited set of Markdown is supported)