Feat: Expose enum types as actual enums in py-rattler

[iamthebot]

There are types like PyPathType that currently represent enums in rust but we expose them with an odd interface where the end-user needs to check for truthiness to ascertain the current value of the enum. Wouldn't it be cleaner to expose these on the python side as enum types?

[baszalmstra]

Do you have an example and how you would change it?

[iamthebot]

For example, the current PathType class in rattler (eg; here) has getters for hardlink, softlink, etc. that just return True if the PathType is of that value. In other words, to get the path type you have to do a bunch of if/else until you hit a possible value that returns True. That's not really idiomatic and a major hassle if you have a lot of values. It also means downstream code needs to worry about any new enum values not being handled.

So for example, to determine the type of a path in Python I have to do something like:

(currently)

...
if path_type.directory:
  return "directory"
elif path_type.softlink:
  return "softlink"
elif path_type.hardlink:
  return "hardlink"

In the rust code in rattler-conda-types however (eg; here this is an enum. That is, the values are mutually exclusive. This means it would be also safe to expose these types as python enums.

If we did that, now it would be easy in python to:

1. Get all possible valid values for the enum
2. Convert from a string to the enum type easily
3. Convert from the enum type to a string
4. Check equality across enums of the same type

For example:

>>> print(path_type.value)
PathType.hardlink

>>> PathType["hardlink"] == PathType.hardlink
`True`

>>> assert path_type == PathType.symlink, "some assertion error"

Doing some cursory research, PyO3 does not natively have a way to export rust enums as Python enums. But a few approaches are suggested in this SO post. Another approach I don't see listed there that might be cleaner (and backwards compatible with older versions of Python) is to expose an int member for these types and on the python side have a matching int based enum with those values. That's more memory efficient too.

[baszalmstra]

Thanks for the writeup! That makes total sense to me. Would you be able to whip up a Pr?

[iamthebot]

Yeah, I can put up a PR Friday.

Update: Ended up a little busy today (4/26) so didn't get to this but will add it next week.

[baszalmstra]

I have been looking at how polars does this. Essentially what they do is that they define an enum in python as a Literal. These are then converted on the Rust side back to enums.

This approach is relatively lightweight on the Python side which I kinda like, but doesn't introduce an "actual type" either which may or may not be nice.

@iamthebot wdyt?

[iamthebot]

@baszalmstra sorry, been ridiculously busy so haven't gotten back to this issue. FWIW there's an open issue in PyO3 on how to best handle enums. Until then...

I'm not a super big fan of the literal option tbh. You still don't get a proper python enum. In that case it's really no better than just using strings. And you miss out on all the magic you get from enums (converting to string representation and back, static analyzers like pyright catching impossible enum values, etc.)

Then again, the least performant option is similar to what we have today in that checking which value you actually in Python involves a ladder of if/else checks on each possible value.

A more performant approach would be to expose the numeric value of the enum to Python (i.e mapping rust enum value -> int inside rust) and inside the python class map that "back" to a proper python enum. But this has the downside of needing to keep the list of possible values in sync in both places. I think we can "cheat" and have a function in rust that returns all the possible enum values as strings. Then python-side we can use Enum's functional API to construct the enum class with those values. That negates us needing to keep the enum values in sync in the python code (since they come from rust).

[baszalmstra]

Mypy does still do type checking when using literals. It will complain when you assign a random string "bla" to a literal. But Im sure mypy can do more magic when using "proper" enums.

With the literal type we do also need to keep the strings in sync between Rust and python though. Unless we generate the enums statically at "compile-time" (whatever that means in python terms) I see no way around that. I realize we could build the type dynamically at runtime, but that will make static analysis like type checking and documentation generation harder.

[iamthebot]

@baszalmstra what would be the downside of having a fn that returns all the enum strings to Python and at import time we construct the corresponding python enum type? Nothing to keep in sync in that case and we get a proper python enum. We could pregenerate stubs for static analysis tools (pyright/mypy/pyre-check)

[baszalmstra]

That would be fine with me as long as we have proper stubs to support documentation generation, IDE autocomplete, and static type checking, otherwise the values will be "invisible" for the user.

We would be happy to update all the enums in the entire codebase but a PR that shows what you have in mind would go a long way!

[Wackyator]

I would love to get started on this issue, @iamthebot do you think you could showcase a simple example for the same?

[iamthebot]

@Wackyator yeah let me put something together today. Sorry, would have started on this earlier but been super busy this week.

[iamthebot]

OK @Wackyator sorry for the delay, was doing some testing on my side to see what would:

1. Work with type checkers (using pyright since that's what we're using here but also tested mypy)
2. Prevent the likelihood of eg; updating enum values in Rust but forgetting to do that in the python package
3. Minimize duplication of enum values as much as possible (since ideally we want to do this for all enums)

It seems that explicitly defining the enum values in Python is going to be the most compatible approach. Then using a unit test to verify that python enums match those in rust. We can send strings from rust to Python which we then convert into the corresponding python enum.

On the rust side (not in py-rattler but in the actual rattler type) the easiest thing would be to eg; adopt something like the strum_macros crate which lets you stringify an enum. Then in the py-rattler rust code for the enums eg; here we just return the string representation of the PrefixPathType enum. Something like a single enum_value method that returns a string instead of the individual boolean methods for is_hardlink, etc. This avoids the need to maintain an extra mapping inside the py-rattler rust code.

Then on the python side..

For example, for PrefixPathType

(note: this uses StrEnum which is Python 3.11+ but a backport package exists for Python 3.6+)

from enum import StrEnum
class PrefixPathType(StrEnum):
    hardlink = "hardlink"
    directory = "directory"
    pyc_file = "pyc_file"
    softlink = "softlink"
    windows_python_entrypoint_script = "windows_python_entrypoint_script"
    windows_python_entrypoint_exe = "windows_python_entrypoint_exe"
    unix_python_entrypoint = "unix_python_entrypoint"
    @classmethod
    def from_py_path_type(cls, py_path_type: PyPrefixPathType) -> PrefixPathType:
        """Construct Rattler PathType from FFI PyPathType object."""
        return PrefixPathType[py_path_type.enum_value()] 

Then in a unit test you can just check that you can check the enum can be constructed for all values of the rust enum.