Skip to content

Commit

Permalink
Add from_bytes classmethod to orm.SinglefileData
Browse files Browse the repository at this point in the history
  • Loading branch information
GeigerJ2 committed Dec 11, 2024
1 parent baf8d7c commit 435e274
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 1 deletion.
11 changes: 11 additions & 0 deletions src/aiida/orm/nodes/data/singlefile.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,17 @@ def from_string(cls, content: str, filename: str | pathlib.Path | None = None, *
"""
return cls(io.StringIO(content), filename, **kwargs)

@classmethod
def from_bytes(
cls, content: bytes, filename: str | pathlib.Path | None = None, **kwargs: t.Any
) -> 'SinglefileData':
"""Construct a new instance and set ``content`` as its contents.
:param content: The content as bytes.
:param filename: Specify filename to use (defaults to ``file.txt``).
"""
return cls(io.BytesIO(content), filename, **kwargs)

def __init__(
self, file: str | pathlib.Path | t.IO, filename: str | pathlib.Path | None = None, **kwargs: t.Any
) -> None:
Expand Down
24 changes: 23 additions & 1 deletion tests/orm/nodes/data/test_singlefile.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,18 +189,40 @@ def test_from_string():
content = 'some\ncontent'
node = SinglefileData.from_string(content).store()
assert node.get_content() == content
assert node.get_content('rb') == content.encode()
assert node.filename == SinglefileData.DEFAULT_FILENAME

content = 'some\ncontent'
filename = 'custom_filename.dat'
node = SinglefileData.from_string(content, filename).store()
assert node.get_content() == content
assert node.get_content('rb') == content.encode()
assert node.filename == filename


def test_from_bytes():
"""Test the :meth:`aiida.orm.nodes.data.singlefile.SinglefileData.from_bytes` classmethod."""
content = b'some\ncontent'
node = SinglefileData.from_bytes(content).store()
assert node.get_content(mode='rb') == content
assert node.get_content(mode='r') == content.decode('utf-8')
assert node.filename == SinglefileData.DEFAULT_FILENAME

content = b'some\ncontent'
filename = 'custom_filename.dat'
node = SinglefileData.from_bytes(content, filename).store()
assert node.get_content(mode='rb') == content
assert node.get_content(mode='r') == content.decode('utf-8')
assert node.filename == filename


def test_get_content():
"""Test the :meth:`aiida.orm.nodes.data.singlefile.SinglefileData.get_content` method."""
content = 'some\ncontent'
node = SinglefileData.from_string(content).store()
assert node.get_content() == content
assert node.get_content('rb') == content.encode()
content = b'some\ncontent'
node = SinglefileData.from_string(content.decode('utf-8')).store()
node = SinglefileData.from_bytes(content).store()
assert node.get_content() == content.decode('utf-8')
assert node.get_content('rb') == content

0 comments on commit 435e274

Please sign in to comment.