Update @particle_input to handle typing that utilzes Union[], List[], etc.

[rocco8773]

@particle_input ques off of arguments annotated with Particle to properly condition that argument into a Particle() object. However, that argument is allowed to be a string (e.g. "He+") or a Particle() object and, thus, should really be annotated with Union[str, Particle]. Without this annotation, IDEs will generally throw warnings if you trying to use "He+" as an argument.

I think @particle_input can accomplish this by doing something like...

def find_particle(func):
     def find_in_depth(args):
         for arg in args:
             if hasattr(arg, "__args__"):
                 result = find_in_depth(arg.__args__)
             else:
                 result = Particle == arg
             
             if result:
                 break
 
         return result
 
     annos = func.__annotations__
     found = {}
 
     for name, anno in annos.items():
         if hasattr(anno, "__args__"):
             results = find_in_depth(anno.__args__)
         else:
             results = Particle == anno

         found[name] = results

     return found

Then the following would result in...

>>> def foo(ion: Union[str, Particle], par: Particle, p2: List[Particle], arg: [u.cm, u.kg]):
...     pass
>>> find_particle(foo)
{'ion': True, 'par': True, 'p2': True, 'arg': False}

[namurphy]

Hm...I wonder if it would be helpful to define particle_like as Union[Particle, str, numbers.Integral], and include that in the same file as @particle_input so that it can be imported in the rest of the package and used for these annotations. That'd probably be cleaner than having Union[..., ..., ...] when reading code and still be understandable.

There's also a new feature in Python 3.9 (which was just released!) for type hinting generics which might be cleaner eventually (e.g., list[str] instead of List[str]). It's too early to use this for now, though.

[namurphy]

Oh...and Particle also accepts integers...which I just realized means that we can do this:

>>> Particle(True)
Particle("H")

Huh.

[rocco8773]

What's False then? Weird!

[namurphy]

Particle(False) raises an InvalidParticleError, instead of the antiparticle of Particle(True). Don't know if we need to have a check for this, since Python itself allows operations like True + 1 which gives 2, and NumPy lets us do numpy.isclose(True, False, True - False) which gives True.

[rocco8773]

I'm currently working on this and hoping to open a PR soon on it. Here are my thoughts on the functionality.

1. We create an typing alias that looks like...

   ParticleLike = typing.Union[str, numbers.Integer, AbstractParticle]
   

   Note that typing.Optional[ParticleLike] == typing.Union[str, numbers.Integer. AbstractParticle, type(None)].

2. @particle_input would then work with typing that look like...

   ParticleLike
   Particle  # any subclass of AbstractParticle
   typing.Optional[ParticleLike]
   typing.Optional[Particle]
   List[Particle]
   List[ParticleLike]
   list(ParticleLike)  # the decorator basically remaps list(), [] -> typing.List[]
   Tuple[ParticleLike, Particle, CustomParticle]
   Tuple[ParticleLike, ...]
   tuple(ParticleLike, ParticleLike, Particle)  # the decorator basically remaps tuple(), () -> typing.Tuple[]
   

3. @particle_input would not work with...

   Dict[str, ParticleLike]
   typing.Union[float, ParticleLike]
   typing.Tuple[typing.Any, ParticleLike]
   typing.List[type.Union[ParticleLike, typing.List[ParticleLike]]]
   typing.Mapping[str, ParticleLike]
   

   In the future we can we can allow for more complex typing variations, but to start I think we should stick to simple ones like ParticleLike, List[ParticleLike], Tuple[ParticleLike], and Optional[ParticleLike].

4. The decorator would raise an Exception under the following conditions...

   1. If ParticleLike is found at a depth of > 1 (annotations are recursively search for ParticleLike).
   2. If a tuple argument does not match the length of the tuple annotation.
   3. If ParticleLike is Union-ed with non-ParticleLike types

5. If an argument is not an instance of one of the AbstractParticle sub-classes, then Particle() will be used to convert it to a particle. This allows CustomParticle and DimensionlessParticle to pass but non-particle objects will only be converted with Particle().

6. I'm not a big fan of the decorator pulling Z and mass_numb from from the particle object. This feels beyond the scope of the decorator, but I will leave it in place for now. I have an idea of how to replace this, which I've outlined in issue #918.

[namurphy]

This looks great, and very thorough! This covers a lot of use cases that I hadn't thought about.

I was thinking that the docstring for this would be a good place to discuss what valid representations of particles are. If you don't mind, I'd like to help out with at least the docstring for this since this is something I've been meaning to put in the documentation anyway.

[rocco8773]

Sure, I've made a bunch of progress in getting the above implement and should be able to open a PR sometime next week. We can go through the docstrings, and more, at that point.

[namurphy]

Sounds good! I also commented on this in #796. In particular, I tentatively created a one-liner for particle_like and started using it, so we could expand off of that.

[rocco8773]

I'll point out, one of the benefits of having a particle_like typing alias is that we will actually have a documentation page that will always be linked back to when we document with `particle_like`.

[namurphy]

EXACTLY!