Skip to content

Commit

Permalink
SinglefileData: Add mode keyword to get_content
Browse files Browse the repository at this point in the history
This allows a user to retrieve the content in bytes. Currently, a user
is forced to use the more elaborate form:

    with singlefile.open(mode='rb') as handle:
        content = handle.read()

or go directly through the repository interface which is a bit hidden
and requires to redundantly specify the filename:

    content = singlefile.base.repository.get_object_content(
	singlefile.filename,
	mode='rb'
    )

these variants can now be simplified to:

    content = singlefile.get_content('rb')
  • Loading branch information
sphuber committed Aug 14, 2023
1 parent b9d087d commit d082df7
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 3 deletions.
7 changes: 4 additions & 3 deletions aiida/orm/nodes/data/singlefile.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,12 +92,13 @@ def open(self, path: str | None = None, mode: t.Literal['r', 'rb'] = 'r') -> t.I
with self.base.repository.open(path, mode=mode) as handle:
yield handle

def get_content(self) -> str:
def get_content(self, mode: str = 'r') -> str | bytes:
"""Return the content of the single file stored for this data node.
:return: the content of the file as a string
:param mode: the mode with which to open the file handle (default: read mode)
:return: the content of the file as a string or bytes, depending on ``mode``.
"""
with self.open(mode='r') as handle: # type: ignore[call-overload]
with self.open(mode=mode) as handle: # type: ignore[call-overload]
return handle.read()

def set_file(self, file: str | t.IO, filename: str | pathlib.Path | None = None) -> None:
Expand Down
8 changes: 8 additions & 0 deletions tests/orm/nodes/data/test_singlefile.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,3 +198,11 @@ def test_from_string():
node = SinglefileData.from_string(content, filename).store()
assert node.get_content() == content
assert node.filename == filename


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

0 comments on commit d082df7

Please sign in to comment.