equinox icon indicating copy to clipboard operation
equinox copied to clipboard

Abstract{Class}Var does not allow string annotations

Open danielward27 opened this issue 1 year ago • 4 comments

This also becomes an issue with from __futures__ import annotations for example

from __future__ import annotations

import equinox as eqx
from equinox import AbstractVar

class AbstractFoo(eqx.Module):
    my_int: AbstractVar[int]

danielward27 avatar Nov 15 '23 15:11 danielward27

I'm afraid this is intentional. String-based annotations have a lot of issues in various edge cases. (And likewise it is best practice not to use from __future__ import annotations.)

Thankfully the whole thing is being improved with PEP 649; take a look at that PEP, and its linked discourse thread, for more details.

patrick-kidger avatar Nov 15 '23 16:11 patrick-kidger

Thanks for the information. Good to hear things will be improving! Right now sphinx requires from __futures__ import annotations to be used in order to have type aliases displayed nicely i.e. ArrayLike rather than typing.Union[jax.Array, numpy.ndarray, numpy.bool_, numpy.number, bool, int, float, complex]. I'm sure I'll find a workaround for now though.

danielward27 avatar Nov 15 '23 17:11 danielward27

My usual trick here is to do an awful hack like:

# in the doc setup script
import some_stdlib_module
some_stdlib_module.GENERATING_DOCUMENTATION = True

# in the library
class HasRepr:
    def __init__(self, string):
        self.string = string

    def __repr__(self):
        return self.string

if getattr(some_stdlib_module, "GENERATING_DOCUMENTATION", False):
    ArrayLike = HasRepr("ArrayLike")
else:
    from jaxtyping import ArrayLike

Side note, having used both, I'd really recommend MkDocs and mkdocs-material over Sphinx. The former means you can just format everything in markdown, which is usually more familiar. Sphinx and rST is usually more awkard, consistently leading to typos due to its various divergences from markdown (e.g. accidentally only using markdown-style single-backticks instead of rST-style double-backticks when wanting unformatted text).

patrick-kidger avatar Nov 15 '23 18:11 patrick-kidger

Great, thanks, that will do for the time being! Yes, I probably should switch, I have like you found sphinx to be on the clunkier side. It's just mostly my laziness when sphinx (mostly) works well and has some convenient extensions.

danielward27 avatar Nov 16 '23 13:11 danielward27