Skip to content

Commit

Permalink
support writing structures to compressed JSON files
Browse files Browse the repository at this point in the history
.json
.json.gz
.json.bz2
.json.xz
.json.lzma
  • Loading branch information
janosh committed May 20, 2023
1 parent 61c20b1 commit ffae4ea
Showing 1 changed file with 13 additions and 13 deletions.
26 changes: 13 additions & 13 deletions pymatgen/core/structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -2493,19 +2493,19 @@ def to(self, filename: str = "", fmt: str = "", **kwargs) -> str | None:
from pymatgen.io.cssr import Cssr

writer = Cssr(self) # type: ignore
elif fmt == "json" or fnmatch(filename.lower(), "*.json"):
s = json.dumps(self.as_dict())
elif fmt == "json" or fnmatch(filename.lower(), "*.json*"):
dct = json.dumps(self.as_dict())
if filename:
with zopen(filename, "wt") as f:
f.write(s)
return s
with zopen(filename, "wt") as file:
file.write(dct)
return dct
elif fmt == "xsf" or fnmatch(filename.lower(), "*.xsf*"):
from pymatgen.io.xcrysden import XSF

s = XSF(self).to_string()
if filename:
with zopen(filename, "wt", encoding="utf8") as f:
f.write(s)
with zopen(filename, "wt", encoding="utf8") as file:
file.write(s)
return s
elif (
fmt == "mcsqs"
Expand All @@ -2517,8 +2517,8 @@ def to(self, filename: str = "", fmt: str = "", **kwargs) -> str | None:

s = Mcsqs(self).to_string()
if filename:
with zopen(filename, "wt", encoding="ascii") as f:
f.write(s)
with zopen(filename, "wt", encoding="ascii") as file:
file.write(s)
return s
elif fmt == "prismatic" or fnmatch(filename, "*prismatic*"):
from pymatgen.io.prismatic import Prismatic
Expand All @@ -2528,8 +2528,8 @@ def to(self, filename: str = "", fmt: str = "", **kwargs) -> str | None:
elif fmt == "yaml" or fnmatch(filename, "*.yaml*") or fnmatch(filename, "*.yml*"):
yaml = YAML()
if filename:
with zopen(filename, "wt") as f:
yaml.dump(self.as_dict(), f)
with zopen(filename, "wt") as file:
yaml.dump(self.as_dict(), file)
return None
sio = StringIO()
yaml.dump(self.as_dict(), sio)
Expand All @@ -2543,8 +2543,8 @@ def to(self, filename: str = "", fmt: str = "", **kwargs) -> str | None:

s = ResIO.structure_to_str(self)
if filename:
with zopen(filename, "wt", encoding="utf8") as f:
f.write(s)
with zopen(filename, "wt", encoding="utf8") as file:
file.write(s)
return None
return s
else:
Expand Down

0 comments on commit ffae4ea

Please sign in to comment.