ENH: Make skops.io.dump and load work with TextIOWrapper #234
Conversation
(note: load already worked, only dump needed change) The idea here is to make it possible to use dump and load with a file wrapper, i.e. using 'with open(...)'. This makes it easier to search and replace pickle dump and load by skops dump and load. Implementation To decide how to check the way to treat the file argument, I went with hasattr(file, "write"), as this is what np.save does. I assume this is more robust than checking the type. Unfortunately, mypy doesn't understand this, so I had to tell it to shut up. I didn't update the docstring or any examples because I don't think we need to 'advertise' this usage.
|
@skops-dev/maintainers ready for review |
skops/io/_persist.py
Outdated
| @@ -39,7 +39,7 @@ def _save(obj: Any) -> io.BytesIO: | |||
| return buffer | |||
|
|
|||
|
|
|||
| def dump(obj: Any, file: str) -> None: | |||
| def dump(obj: Any, file: str | io.TextIOWrapper) -> None: | |||
There was a problem hiding this comment.
is it a textiowrapper or anything which has write?
There was a problem hiding this comment.
It's the latter, if you prefer I can create a protocol for objects with write attributes, just thought TextIOWrapper is the most common case by far.
There was a problem hiding this comment.
Or not have it typed? 👻
We also accept a Path.
What we really accept here, is a file-like object, or a path like object. And that's best explained in the docstring w/o any types.
There was a problem hiding this comment.
I dug a bit further, it seems that typing.BinaryIO would be the correct type:
https://docs.python.org/3/library/typing.html#typing.BinaryIO
This allows mypy to flag code like the one below which uses the wrong mode:
with open(file, 'w') as f: # <= should be 'wb' mode
dump(f)WDYT?
Regarding Path: Yeah, that's supported, I originally didn't add it because it's not related to this PR, but I can still add it here.
There was a problem hiding this comment.
Ok, I could live with BinaryIO then.
There was a problem hiding this comment.
Okay, I changed it accordingly. I also changed to if isinstance(file, (str, Path)), which is now correctly type-checked by mypy.
BenjaminBossan
changed the title
BenjaminBossan
changed the title
Done, can be reviewed again. |
Co-authored-by: Adrin Jalali <adrin.jalali@gmail.com>
(note:
loadalready worked, onlydumpneeded change)The idea here is to make it possible to use
skops.io.dumpandloadwith a file wrapper, i.e. usingwith open(...). This makes it easier to search and replacepickle.dumpandloadbyskops.io.dumpandload.Implementation
To decide how to check the way to treat the file argument, I went with
hasattr(file, "write"), as this is whatnp.savedoes. I assume this is more robust than checking the type. Unfortunately, mypy doesn't understand this, so I had to tell it to shut up.I didn't update the docstring or any examples because I don't think we need to "advertise" this usage.