Adjust for pandas < 2.0

This commit is contained in:
Alex Burlacu 2023-08-11 19:56:08 +03:00
parent d4b11dfa22
commit 83a04d438c

View File

@ -1,4 +1,5 @@
import gzip import gzip
import io
import json import json
import yaml import yaml
import mimetypes import mimetypes
@ -1062,7 +1063,16 @@ class Artifacts(object):
# (otherwise it is encoded and creates new hash every time) # (otherwise it is encoded and creates new hash every time)
if self._compression == "gzip": if self._compression == "gzip":
with gzip.GzipFile(local_filename, 'wb', mtime=0) as gzip_file: with gzip.GzipFile(local_filename, 'wb', mtime=0) as gzip_file:
artifact_object.to_csv(gzip_file, **kwargs) try:
pd_version = int(pd.__version__.split(".")[0])
except ValueError:
pd_version = 0
if pd_version >= 2:
artifact_object.to_csv(gzip_file, **kwargs)
else:
# old (pandas<2) versions of pandas cannot handle direct gzip stream, so we manually encode it
artifact_object.to_csv(io.TextIOWrapper(gzip_file), **kwargs)
else: else:
artifact_object.to_csv(local_filename, compression=self._compression) artifact_object.to_csv(local_filename, compression=self._compression)