Skip to content
This repository was archived by the owner on Aug 25, 2024. It is now read-only.

Commit

Permalink
source: json: Load JSON and patch label in dump
Browse files Browse the repository at this point in the history
  • Loading branch information
sudharsana-kjl authored and pdxjohnny committed Jun 14, 2019
1 parent 5b3a970 commit 2302fb8
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 10 deletions.
2 changes: 1 addition & 1 deletion dffml/source/file.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ async def _close(self):
elif self.config.filename[::-1].startswith((".zip")[::-1]):
close = self.zip_closer_helper()
else:
close = open(self.config.filename, "w")
close = open(self.config.filename, "w+")
with close as fd:
await self.dump_fd(fd)

Expand Down
19 changes: 11 additions & 8 deletions dffml/source/json.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,21 +17,24 @@ class JSONSource(FileSource, MemorySource):
stored in memory.
"""

def __init__(self, config):
super().__init__(config)
self.repos = {}

async def load_fd(self, fd):
self.repos = json.load(fd)
repos = json.load(fd)
self.mem = {
src_url: Repo(src_url, data=data)
for src_url, data in self.repos.get(self.config.label, {}).items()
for src_url, data in repos.get(self.config.label, {}).items()
}
LOGGER.debug("%r loaded %d records", self, len(self.mem))

async def dump_fd(self, fd):
self.repos[self.config.label] = {
repos = {}
if fd.seekable():
# Empty Stream is not a Valid JSON
if fd.seek(0) is not 0:
repos = json.load(fd)
fd.seek(0)
fd.truncate(0)
repos[self.config.label] = {
repo.src_url: repo.dict() for repo in self.mem.values()
}
json.dump(self.repos, fd)
json.dump(repos, fd)
LOGGER.debug("%r saved %d records", self, len(self.mem))
2 changes: 1 addition & 1 deletion tests/source/test_file.py
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ async def test_close(self):
):
async with FakeFileSource(self.config("testfile")):
pass
m_open.assert_called_once_with("testfile", "w")
m_open.assert_called_once_with("testfile", "w+")

async def test_close_gz(self):
m_open = mock_open()
Expand Down

0 comments on commit 2302fb8

Please sign in to comment.