[Sugar] Add `PythonicPytree.replace()`

[georgematheos]

In some of my code, I've been subclassing the following rather than PythonicPytree:

from dataclasses import replace

class UpdatablePytree(PythonicPytree):
    """A Pytree supporting self.replace(...) to functionally update fields."""

    def replace(self, **kwargs):
        return replace(self, **kwargs)

Any myobj: MyClass where MyClass is a sub-type of UpdateablePytree supports syntax like

mynewobj = myobj.replace(x=new_x_value, y=new_y_value)

to functionally update the values on the x and y field in myobj.

This issue proposes to add this implementation of .replace to PythonicPytree directly.

[linear]

GEN-857 [Sugar] Add `PythonicPytree.replace()`

In some of my code, I've been subclassing the following rather than PythonicPytree:

class UpdatablePytree(PythonicPytree):
    """A Pytree supporting self.replace(...) to functionally update fields."""
 
    def replace(self, **kwargs):
        return replace(self, **kwargs)

Any myobj: MyClass where MyClass is a sub-type of UpdateablePytree supports syntax like

mynewobj = myobj.replace(x=new_x_value, y=new_y_value)

This issue proposes to add this implementation of .replace to PythonicPytree directly.

[sritchie]

Wdyt of sending this as a PR, @georgematheos ? Looks good to me!

[georgematheos]

Can do!

I took a look at making a PR, but realized it is a little more complicated than I had expected, since replace only works on dataclasses, whereas Pytree and PythonicPytrees need not be dataclasses. So I get a type warning/error if I try to add this replace function to the Pytree of PythonicPytree class.

I can come back to this later, but need to switch back to my other work for now!

[sritchie]

@georgematheos I think they actually do need to be dataclasses, though the bottom-most subclass has to be annotated as a dataclass. Pytree has an __init__ that will check for this. But I believe you about the type error, push what you've got when you get a minute and we can try it out.