Teach unxt about your own type

Teach unxt about your own type#

Take a class that unxt has never heard of. By the end of this page, unxt’s own functions will work on it — unit_of, dimension_of, ustrip, uconvert — with nothing subclassed and not a line of unxt touched. A function written against unxt’s API will accept our class as readily as it accepts a Quantity.

You need unxt installed and nothing else. (unxts.api comes with it.)

Set up#

>>> from plum import dispatch
>>> import unxt as u

Start with a type unxt does not know#

Here is a small class standing in for whatever your project already has — a raw instrument reading, a number with a unit name attached:

>>> class Reading:
...     """A raw instrument reading: a number and a unit name."""
...     def __init__(self, value, unit):
...         self.value = value
...         self.unit = unit
...     def __repr__(self):
...         return f"Reading({self.value}, {self.unit!r})"

>>> r = Reading(2000.0, "m")

Ask unxt about it:

>>> print(u.unit_of(r))
None

None — not an error. unit_of has a fallback for objects it does not recognise, and as far as it is concerned our Reading is just some object. The "m" sitting right there in the attribute means nothing to it yet.

Register the first implementation#

Now we tell unxt how to find the unit. Define a function named unit_of, annotate the argument with our type, and decorate it with plum’s dispatch:

>>> @dispatch
... def unit_of(obj: Reading, /):
...     return u.unit(obj.unit)

That is the entire registration. Ask again:

>>> u.unit_of(r)
Unit("m")

Look carefully at what just happened. We called u.unit_of — unxt’s function, the one we did not modify — and it dispatched into the implementation we wrote a moment ago:

>>> unit_of is u.unit_of
True

The same object, with one more method on it. We never imported unit_of, and we did not need to: plum.dispatch is a single shared dispatcher, and it files module-level functions under their bare __name__. Writing @dispatch on a function called unit_of puts your implementation into the one and only unit_of, wherever in your codebase you wrote it.

That cuts both ways, so it is worth saying plainly: the name is the whole key. A @dispatch def unit_of you meant as your own private helper would join unxt’s function too. Name these after the API function you intend to extend, and nothing else.

Add the dimension#

Same pattern. Once unxt can get the unit, the dimension follows from it:

>>> @dispatch
... def dimension_of(obj: Reading, /):
...     return u.dimension_of(u.unit_of(obj))

>>> u.dimension_of(r)
PhysicalType('length')

Notice our implementation called u.unit_of(obj) — the function we had just extended. Registrations compose; each one you add makes the next one shorter.

Watch generic code accept it#

Here is the payoff. Write a function that knows nothing about Reading and only uses unxt’s API:

>>> def describe(x):
...     return f"{u.dimension_of(x)} measured in {u.unit_of(x)}"

Hand it a Quantity:

>>> describe(u.Q(5.0, "m"))
'length measured in m'

and hand it our Reading:

>>> describe(r)
'length measured in m'

The same function, unchanged, for both. describe could have been written by someone else, in another library, before our class existed — this is what the abstract API is for.

Make it convert#

Two more registrations and Reading supports unit conversion. ustrip returns the bare number in the requested unit:

>>> @dispatch
... def ustrip(to: u.AbstractUnit, obj: Reading, /):
...     return u.ustrip(to, u.Q(obj.value, obj.unit))

>>> u.ustrip(u.unit("km"), r)
Array(2., dtype=float32, weak_type=True)

and uconvert returns a converted object of our own type:

>>> @dispatch
... def uconvert(to: u.AbstractUnit, obj: Reading, /):
...     return Reading(float(u.ustrip(to, obj)), str(to))

>>> u.uconvert(u.unit("km"), r)
Reading(2.0, 'km')

2000 metres came back as 2.0 kilometres, still a Reading. We borrowed unxt’s conversion machinery for the arithmetic and kept our own type on the outside.

What we built#

Reading now answers to four of unxt’s functions, and generic unit-aware code works on it. We never subclassed an unxt class, never edited unxt, and never imported jax — which is the point of unxts.api being a separate package: a library can speak this API while leaving the choice of whether to install unxt itself to its users.

Where to go next#