Skip to content

Commit

Permalink
[BugFix] extract gz into target dir (#3389)
Browse files Browse the repository at this point in the history
  • Loading branch information
Rhett-Ying committed Sep 30, 2021
1 parent e234fcf commit f9fd7fd
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 1 deletion.
3 changes: 2 additions & 1 deletion python/dgl/data/utils.py
Expand Up @@ -219,7 +219,8 @@ def extract_archive(file, target_dir, overwrite=False):
import gzip
import shutil
with gzip.open(file, 'rb') as f_in:
with open(file[:-3], 'wb') as f_out:
target_file = os.path.join(target_dir, os.path.basename(file)[:-3])
with open(target_file, 'wb') as f_out:
shutil.copyfileobj(f_in, f_out)
elif file.endswith('.zip'):
import zipfile
Expand Down
18 changes: 18 additions & 0 deletions tests/compute/test_data.py
Expand Up @@ -2,6 +2,9 @@
import unittest
import backend as F
import numpy as np
import gzip
import tempfile
import os


@unittest.skipIf(F._default_context_str == 'gpu', reason="Datasets don't need to be tested on GPU.")
Expand Down Expand Up @@ -139,10 +142,25 @@ def test_reddit():
assert np.array_equal(dst, np.sort(dst))


@unittest.skipIf(F._default_context_str == 'gpu', reason="Datasets don't need to be tested on GPU.")
def test_extract_archive():
# gzip
with tempfile.TemporaryDirectory() as src_dir:
gz_file = 'gz_archive'
gz_path = os.path.join(src_dir, gz_file + '.gz')
content = b"test extract archive gzip"
with gzip.open(gz_path, 'wb') as f:
f.write(content)
with tempfile.TemporaryDirectory() as dst_dir:
data.utils.extract_archive(gz_path, dst_dir, overwrite=True)
assert os.path.exists(os.path.join(dst_dir, gz_file))


if __name__ == '__main__':
test_minigc()
test_gin()
test_data_hash()
test_tudataset_regression()
test_fraud()
test_fakenews()
test_extract_archive()

0 comments on commit f9fd7fd

Please sign in to comment.