diff --git a/PyOpenWorm/bittorrent.py b/PyOpenWorm/bittorrent.py index cdaca3604..b0fff342d 100644 --- a/PyOpenWorm/bittorrent.py +++ b/PyOpenWorm/bittorrent.py @@ -8,6 +8,7 @@ from PyOpenWorm.datasource_loader import DataSourceDirLoader import pickle import os.path +from os.path import join as p from googleapiclient.discovery import build from google_auth_oauthlib.flow import InstalledAppFlow from google.auth.transport.requests import Request @@ -16,7 +17,7 @@ # If modifying these scopes, delete the file token.pickle. SCOPES = ['https://www.googleapis.com/auth/drive'] -def download_torrent(torrent_name): +def download_torrent(torrent_name, powdir='.pow', tempdir='.'): """Shows basic usage of the Drive v3 API. Prints the names and ids of the first 10 files the user has access to. """ @@ -24,8 +25,8 @@ def download_torrent(torrent_name): # The file token.pickle stores the user's access and refresh tokens, and is # created automatically when the authorization flow completes for the first # time. - if os.path.exists('.pow/temp/token.pickle'): - with open('.pow/temp/token.pickle', 'rb') as token: + if os.path.exists(p(powdir, 'temp', 'token.pickle')): + with open(p(powdir, 'temp', 'token.pickle'), 'rb') as token: creds = pickle.load(token) # If there are no (valid) credentials available, let the user log in. if not creds or not creds.valid: @@ -33,10 +34,10 @@ def download_torrent(torrent_name): creds.refresh(Request()) else: flow = InstalledAppFlow.from_client_secrets_file( - '.pow/temp/credentials.json', SCOPES) + p(powdir, tempdir, 'credentials.json'), SCOPES) creds = flow.run_local_server() # Save the credentials for the next run - with open('.pow/temp/token.pickle', 'wb') as token: + with open(p(powdir, tempdir, 'token.pickle'), 'wb') as token: pickle.dump(creds, token) service = build('drive', 'v3', credentials=creds) @@ -69,23 +70,26 @@ def download_torrent(torrent_name): while done is False: status, done = downloader.next_chunk() print("Download %d%%." % int(status.progress() * 100)) - stri = "./" + selected_file_name + stri = p(tempdir, selected_file_name) with io.open(stri, 'wb') as f: fh.seek(0) f.write(fh.read()) - return selected_file_name + return p(tempdir, selected_file_name) + class BitTorrentDataSourceDirLoader(DataSourceDirLoader): + + def __init__(self, *args, powdir, tempdir, **kwargs): + super(BitTorrentDataSourceDirLoader, self).__init__(*args, **kwargs) + self.tempdir = tempdir + self.powdir = powdir + def load(self, *data_source): for d in data_source: x = list(d.torrent_file_name()) - downloaded_torrent_name = download_torrent(x[0]) + downloaded_torrent_name = download_torrent(x[0], powdir=self.powdir, tempdir=self.tempdir) print('downloaded torrent', downloaded_torrent_name) - os.system("torrent_cli.py start &") - os.system("torrent_cli.py add "+ downloaded_torrent_name) - - - + os.system("torrent_cli.py add " + downloaded_torrent_name) diff --git a/PyOpenWorm/command.py b/PyOpenWorm/command.py index d179ae5bc..9673eef63 100644 --- a/PyOpenWorm/command.py +++ b/PyOpenWorm/command.py @@ -1155,7 +1155,7 @@ def _load_data_source_directories(self): # XXX persist the dict lclasses = [POWDirDataSourceDirLoader(), - BitTorrentDataSourceDirLoader()] + BitTorrentDataSourceDirLoader(powdir=self.powdir)] dsd = _DSD(dict(), pth_join(self.powdir, 'data_source_data'), lclasses) try: dindex = open(pth_join(self.powdir, 'data_source_directories')) diff --git a/tests/BitTorrentTest.py b/tests/BitTorrentTest.py index 4e0b680db..dc0a4f41c 100644 --- a/tests/BitTorrentTest.py +++ b/tests/BitTorrentTest.py @@ -8,6 +8,7 @@ from .DataTestTemplate import _DataTest from PyOpenWorm.bittorrent import BitTorrentDataSourceDirLoader import os +from os.path import exists, join as p import six import sys import tempfile @@ -16,6 +17,7 @@ import transaction from pytest import mark, fixture import unittest +import shutil from PyOpenWorm.data_trans.local_file_ds import LocalFileDataSource as LFDS from PyOpenWorm import connect @@ -27,29 +29,33 @@ class TestBitTorrentDataSourceDirLoader(_DataTest): def setUp(self): super(TestBitTorrentDataSourceDirLoader, self).setUp() + self.ident = 'http://openworm.org/entities/ConnectomeCSVDataSource/Mark_Arnab_3000_connections' with transaction.manager: # Create data sources ctx = Context(ident='http://example.org/context', conf=self.connection.conf) ctx(LFDS)( - ident='http://example.org/lfds', + ident=self.ident, file_name='Merged_Nuclei_Stained_Worm.zip', torrent_file_name='d9da5ce947c6f1c127dfcdc2ede63320.torrent' ) ctx.save_context() + self.tempdir = tempfile.mkdtemp() + def tearDown(self): + shutil.rmtree(self.tempdir) def test_torrent_download1(self): ctx = Context(ident="http://example.org/context", conf=self.connection.conf) - self.assertFalse(os.path.exists("d9da5ce947c6f1c127dfcdc2ede63320.torrent"), False) - self.assertFalse(os.path.exists("Merged_Nuclei_Stained_Worm.zip"), False) + self.assertFalse(exists(p(self.tempdir, "d9da5ce947c6f1c127dfcdc2ede63320.torrent")), False) + self.assertFalse(exists(p(self.tempdir, "Merged_Nuclei_Stained_Worm.zip")), False) - content = BitTorrentDataSourceDirLoader("./") - ident = 'http://openworm.org/entities/ConnectomeCSVDataSource/Mark_Arnab_3000_connections' + content = BitTorrentDataSourceDirLoader(self.tempdir, powdir='.pow', tempdir=self.tempdir) - for m in ctx.stored(LFDS)().load(): + for m in ctx.stored(LFDS)(ident=self.ident).load(): + print("GOT", m) content_path = content.load(m) - self.assertTrue(os.path.exists("d9da5ce947c6f1c127dfcdc2ede63320.torrent"), True) - self.assertTrue(os.path.exists("Merged_Nuclei_Stained_Worm.zip"), True) + self.assertTrue(exists(p(self.tempdir, "d9da5ce947c6f1c127dfcdc2ede63320.torrent")), True) + self.assertTrue(exists(p(self.tempdir, "Merged_Nuclei_Stained_Worm.zip")), True) # Merged_Nuclei_Stained_Worm.zip will appear but its contents take a while to download # watch the progress with - 'watch python3 torrent_cli.py'