Apply PR feedback

libp2p/network/connection/raw_connection.py

@@ -53,5 +53,6 @@ class RawConnection(IRawConnection):
 
     async def close(self) -> None:
         self.writer.close()
-        if sys.version_info[0:2] > (3, 6):
-            await self.writer.wait_closed()
+        if sys.version_info < (3, 7):
+            return
+        await self.writer.wait_closed()

libp2p/tools/factories.py

@@ -1,6 +1,8 @@
 import asyncio
 from typing import Any, AsyncIterator, Dict, Tuple, cast
 
+# NOTE: import ``asynccontextmanager`` from ``contextlib`` when support for python 3.6 is dropped.
+from async_generator import asynccontextmanager
 import factory
 
 from libp2p import generate_new_rsa_identity, generate_peer_id_from
@@ -32,12 +34,6 @@ from .constants import (
 )
 from .utils import connect, connect_swarm
 
-try:
-    from contextlib import asynccontextmanager
-except ImportError:
-    # NOTE: mypy complains about a duplicate import without the following ``# type: ignore``
-    from async_generator import asynccontextmanager  # type: ignore
-
 
 def initialize_peerstore_with_our_keypair(self_id: ID, key_pair: KeyPair) -> PeerStore:
     peer_store = PeerStore()
@@ -177,7 +173,7 @@ async def host_pair_factory(is_secure: bool) -> Tuple[BasicHost, BasicHost]:
     return hosts[0], hosts[1]
 
 
-@asynccontextmanager
+@asynccontextmanager  # type: ignore
 async def pair_of_connected_hosts(
     is_secure: bool = True
 ) -> AsyncIterator[Tuple[BasicHost, BasicHost]]:

libp2p/transport/tcp/tcp.py

@@ -54,9 +54,11 @@ class TCPListener(IListener):
         if self.server is None:
             return
         self.server.close()
-        if sys.version_info[0:2] > (3, 6):
-            await self.server.wait_closed()
+        server = self.server
         self.server = None
+        if sys.version_info < (3, 7):
+            return
+        await server.wait_closed()
 
 
 class TCP(ITransport):