Commit 16af6550 authored by Tiago Peixoto's avatar Tiago Peixoto
Browse files

Add support for zstandard compression

parent 3ad01bce
...@@ -115,14 +115,9 @@ public: ...@@ -115,14 +115,9 @@ public:
std::streamsize write(const char* s, std::streamsize n) std::streamsize write(const char* s, std::streamsize n)
{ {
string buf(s, s+n); string buf(s, s+n);
#if (PY_MAJOR_VERSION >= 3)
// in python 3 we need to construct a 'bytes' instance
PyObject* bytes = PyBytes_FromStringAndSize(s, n); PyObject* bytes = PyBytes_FromStringAndSize(s, n);
boost::python::handle<> x(bytes); boost::python::handle<> x(bytes);
boost::python::object pbuf(x); boost::python::object pbuf(x);
#else
boost::python::str pbuf(buf);
#endif
_file.attr("write")(pbuf); _file.attr("write")(pbuf);
return n; return n;
} }
......
...@@ -117,6 +117,10 @@ try: ...@@ -117,6 +117,10 @@ try:
import lzma import lzma
except ImportError: except ImportError:
pass pass
try:
import zstandard
except ImportError:
pass
import weakref import weakref
import copy import copy
import textwrap import textwrap
...@@ -124,6 +128,7 @@ import io ...@@ -124,6 +128,7 @@ import io
import collections.abc import collections.abc
import itertools import itertools
import csv import csv
import contextlib
from .decorators import _require, _limit_args, _copy_func from .decorators import _require, _limit_args, _copy_func
...@@ -2793,7 +2798,8 @@ class Graph(object): ...@@ -2793,7 +2798,8 @@ class Graph(object):
def __get_file_format(self, file_name): def __get_file_format(self, file_name):
fmt = None fmt = None
for f in ["gt", "graphml", "xml", "dot", "gml"]: for f in ["gt", "graphml", "xml", "dot", "gml"]:
names = ["." + f, ".%s.gz" % f, ".%s.bz2" % f, ".%s.xz" % f] names = ["." + f, ".%s.gz" % f, ".%s.bz2" % f, ".%s.xz" % f,
".%s.zst" % f]
for name in names: for name in names:
if file_name.endswith(name): if file_name.endswith(name):
fmt = f fmt = f
...@@ -2823,16 +2829,29 @@ class Graph(object): ...@@ -2823,16 +2829,29 @@ class Graph(object):
if isinstance(file_name, str): if isinstance(file_name, str):
file_name = os.path.expanduser(file_name) file_name = os.path.expanduser(file_name)
f = open(file_name) # throw the appropriate exception, if not found with open(file_name) as f: # throw the appropriate exception
pass
if fmt == 'auto' and isinstance(file_name, str): if fmt == 'auto' and isinstance(file_name, str):
fmt = self.__get_file_format(file_name) fmt = self.__get_file_format(file_name)
elif fmt == "auto": elif fmt == "auto":
fmt = "gt" fmt = "gt"
if isinstance(file_name, str) and file_name.endswith(".xz"): fctx = contextlib.suppress()
try: if isinstance(file_name, str):
file_name = lzma.open(file_name, mode="rb") if file_name.endswith(".xz"):
except NameError: try:
raise NotImplementedError("lzma compression is only available in Python >= 3.3") fctx = file_name = lzma.open(file_name, mode="rb")
except NameError:
raise NotImplementedError("lzma compression is only available in Python >= 3.3")
elif file_name.endswith(".zst"):
try:
cctx = zstandard.ZstdDecompressor()
except NameError:
raise NotImplementedError("zstandard module not installed")
fctx = contextlib.ExitStack()
f = open(file_name, mode="rb")
fctx.enter_context(f)
file_name = cctx.stream_reader(f)
fctx.enter_context(file_name)
if fmt == "graphml": if fmt == "graphml":
fmt = "xml" fmt = "xml"
if ignore_vp is None: if ignore_vp is None:
...@@ -2846,8 +2865,10 @@ class Graph(object): ...@@ -2846,8 +2865,10 @@ class Graph(object):
fmt, ignore_vp, fmt, ignore_vp,
ignore_ep, ignore_gp) ignore_ep, ignore_gp)
else: else:
props = self.__graph.read_from_file("", file_name, fmt, with fctx:
ignore_vp, ignore_ep, ignore_gp) props = self.__graph.read_from_file("", file_name, fmt,
ignore_vp, ignore_ep,
ignore_gp)
for name, prop in props[0].items(): for name, prop in props[0].items():
self.vertex_properties[name] = VertexPropertyMap(prop, self) self.vertex_properties[name] = VertexPropertyMap(prop, self)
for name, prop in props[1].items(): for name, prop in props[1].items():
...@@ -2908,22 +2929,34 @@ class Graph(object): ...@@ -2908,22 +2929,34 @@ class Graph(object):
if fmt == "graphml": if fmt == "graphml":
fmt = "xml" fmt = "xml"
if isinstance(file_name, str) and file_name.endswith(".xz"): fctx = contextlib.suppress()
try: if isinstance(file_name, str):
file_name = lzma.open(file_name, mode="wb") if file_name.endswith(".xz"):
except NameError: try:
raise NotImplementedError("lzma compression is only available in Python >= 3.3") fctx = file_name = lzma.open(file_name, mode="wb")
except NameError:
raise NotImplementedError("lzma compression is only available in Python >= 3.3")
elif file_name.endswith(".zst"):
try:
cctx = zstandard.ZstdCompressor(level=19)
except NameError:
raise NotImplementedError("zstandard module not installed")
fctx = contextlib.ExitStack()
f = open(file_name, mode="wb")
fctx.enter_context(f)
file_name = cctx.stream_writer(f)
fctx.enter_context(file_name)
props = [(name[1], prop._PropertyMap__map) for name, prop in \ props = [(name[1], prop._PropertyMap__map) for name, prop in \
u.__properties.items()] u.__properties.items()]
if isinstance(file_name, str): if isinstance(file_name, str):
f = open(file_name, "w") # throw the appropriate exception, if with open(file_name, "w"): # throw the appropriate exception, if
# unable to open pass # unable to open
f.close()
u.__graph.write_to_file(file_name, None, fmt, props) u.__graph.write_to_file(file_name, None, fmt, props)
else: else:
u.__graph.write_to_file("", file_name, fmt, props) with fctx:
u.__graph.write_to_file("", file_name, fmt, props)
# Directedness # Directedness
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment