diff --git a/trio_ftplib/__init__.py b/trio_ftplib/__init__.py index 3934f16..4b4f361 100644 --- a/trio_ftplib/__init__.py +++ b/trio_ftplib/__init__.py @@ -1,2 +1,3 @@ from ._aftplib import AFTP from .operations import download_text +from ftplib import Error as FTPLibError diff --git a/trio_ftplib/_aftplib.py b/trio_ftplib/_aftplib.py index 2713a13..d055a67 100644 --- a/trio_ftplib/_aftplib.py +++ b/trio_ftplib/_aftplib.py @@ -58,15 +58,20 @@ class AFTP(ftplib.FTP): def _finalize(): fp.close() conn.close() - self.voidresp() - try: - while True: + while True: + try: line = await trio.to_thread.run_sync(fp.readline) - if not line: - break + except Exception: + await trio.to_thread.run_sync(_finalize) + raise + if not line: + await trio.to_thread.run_sync(_finalize) + await trio.to_thread.run_sync(self.voidresp) # check the response code + break + try: yield line - finally: - await trio.to_thread.run_sync(_finalize) - + except GeneratorExit: + await trio.to_thread.run_sync(_finalize) + raise diff --git a/trio_ftplib/_tests/test_tso.py b/trio_ftplib/_tests/test_tso.py index 443d6f3..6b1c12b 100644 --- a/trio_ftplib/_tests/test_tso.py +++ b/trio_ftplib/_tests/test_tso.py @@ -8,3 +8,4 @@ async def test_tso_sinaf(): path='MIC.NAF.MZ.BBD2.M130D001.D200306')) as _text: async for line in _text: assert line.count('@') == 15 + break