From e5eb222799fd509efdfff4ed7f07e4e2c3902453 Mon Sep 17 00:00:00 2001
From: c056091 <clovis-fabricio.costa@caixa.gov.br>
Date: Fri, 13 Mar 2020 17:07:28 -0300
Subject: [PATCH] finalization

---
 trio_ftplib/__init__.py        |  1 +
 trio_ftplib/_aftplib.py        | 21 +++++++++++++--------
 trio_ftplib/_tests/test_tso.py |  1 +
 3 files changed, 15 insertions(+), 8 deletions(-)

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