Skip to content

Commit 5d47177

Browse files
committed
just call asyncio.current_task() directly
1 parent 77ca1a5 commit 5d47177

File tree

2 files changed

+15
-26
lines changed

2 files changed

+15
-26
lines changed

custom_components/pyscript/handler.py

Lines changed: 3 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,6 @@
1111
_LOGGER = logging.getLogger(LOGGER_PATH + ".handler")
1212

1313

14-
def current_task():
15-
"""Return our asyncio current task."""
16-
try:
17-
# python >= 3.7
18-
return asyncio.current_task()
19-
except AttributeError:
20-
# python <= 3.6
21-
return asyncio.tasks.Task.current_task()
22-
23-
2414
class Handler:
2515
"""Define function handler functions."""
2616

@@ -74,7 +64,7 @@ async def task_unique(self, name, kill_me=False):
7464
"""Implement task.unique()."""
7565
if name in self.unique_name2task:
7666
if kill_me:
77-
task = current_task()
67+
task = asyncio.current_task()
7868

7969
# it seems we need to use another task to cancel ourselves
8070
# I'm sure there is a better way to cancel ourselves...
@@ -97,7 +87,7 @@ async def cancel_self():
9787
await task
9888
except asyncio.CancelledError:
9989
pass
100-
task = current_task()
90+
task = asyncio.current_task()
10191
if task in self.our_tasks:
10292
self.unique_name2task[name] = task
10393
self.unique_task2name[task] = name
@@ -204,7 +194,7 @@ async def run_coro(self, coro):
204194
#
205195
# Add a placeholder for the new task so we know it's one we started
206196
#
207-
task = current_task()
197+
task = asyncio.current_task()
208198
self.our_tasks.add(task)
209199
try:
210200
await coro

custom_components/pyscript/jupyter_kernel.py

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import uuid
2121

2222
from .const import LOGGER_PATH
23-
from .handler import current_task
2423

2524
_LOGGER = logging.getLogger(LOGGER_PATH + ".jupyter_kernel")
2625

@@ -522,7 +521,7 @@ async def control_listen(self, reader, writer):
522521
"""Task that listens to control messages."""
523522
try:
524523
_LOGGER.debug("control_listen connected")
525-
await self.housekeep_q.put(["register", "control", current_task()])
524+
await self.housekeep_q.put(["register", "control", asyncio.current_task()])
526525
control_socket = ZmqSocket(reader, writer, "ROUTER")
527526
await control_socket.handshake()
528527
while 1:
@@ -535,7 +534,7 @@ async def control_listen(self, reader, writer):
535534
raise
536535
except EOFError:
537536
_LOGGER.debug("control_listen got eof")
538-
await self.housekeep_q.put(["unregister", "control", current_task()])
537+
await self.housekeep_q.put(["unregister", "control", asyncio.current_task()])
539538
control_socket.close()
540539
except Exception as err: # pylint: disable=broad-except
541540
_LOGGER.error("control_listen exception %s", err)
@@ -545,7 +544,7 @@ async def stdin_listen(self, reader, writer):
545544
"""Task that listens to stdin messages."""
546545
try:
547546
_LOGGER.debug("stdin_listen connected")
548-
await self.housekeep_q.put(["register", "stdin", current_task()])
547+
await self.housekeep_q.put(["register", "stdin", asyncio.current_task()])
549548
stdin_socket = ZmqSocket(reader, writer, "ROUTER")
550549
await stdin_socket.handshake()
551550
while 1:
@@ -555,7 +554,7 @@ async def stdin_listen(self, reader, writer):
555554
raise
556555
except EOFError:
557556
_LOGGER.debug("stdin_listen got eof")
558-
await self.housekeep_q.put(["unregister", "stdin", current_task()])
557+
await self.housekeep_q.put(["unregister", "stdin", asyncio.current_task()])
559558
stdin_socket.close()
560559
except Exception: # pylint: disable=broad-except
561560
_LOGGER.error("stdin_listen exception %s", traceback.format_exc(-1))
@@ -565,7 +564,7 @@ async def shell_listen(self, reader, writer):
565564
"""Task that listens to shell messages."""
566565
try:
567566
_LOGGER.debug("shell_listen connected")
568-
await self.housekeep_q.put(["register", "shell", current_task()])
567+
await self.housekeep_q.put(["register", "shell", asyncio.current_task()])
569568
shell_socket = ZmqSocket(reader, writer, "ROUTER")
570569
await shell_socket.handshake()
571570
while 1:
@@ -576,7 +575,7 @@ async def shell_listen(self, reader, writer):
576575
raise
577576
except EOFError:
578577
_LOGGER.debug("shell_listen got eof")
579-
await self.housekeep_q.put(["unregister", "shell", current_task()])
578+
await self.housekeep_q.put(["unregister", "shell", asyncio.current_task()])
580579
shell_socket.close()
581580
except Exception: # pylint: disable=broad-except
582581
_LOGGER.error("shell_listen exception %s", traceback.format_exc(-1))
@@ -586,7 +585,7 @@ async def heartbeat_listen(self, reader, writer):
586585
"""Task that listens and responds to heart beat messages."""
587586
try:
588587
_LOGGER.debug("heartbeat_listen connected")
589-
await self.housekeep_q.put(["register", "heartbeat", current_task()])
588+
await self.housekeep_q.put(["register", "heartbeat", asyncio.current_task()])
590589
heartbeat_socket = ZmqSocket(reader, writer, "REP")
591590
await heartbeat_socket.handshake()
592591
while 1:
@@ -597,7 +596,7 @@ async def heartbeat_listen(self, reader, writer):
597596
raise
598597
except EOFError:
599598
_LOGGER.debug("heartbeat_listen got eof")
600-
await self.housekeep_q.put(["unregister", "heartbeat", current_task()])
599+
await self.housekeep_q.put(["unregister", "heartbeat", asyncio.current_task()])
601600
heartbeat_socket.close()
602601
except Exception: # pylint: disable=broad-except
603602
_LOGGER.error("heartbeat_listen exception: %s", traceback.format_exc(-1))
@@ -607,7 +606,7 @@ async def iopub_listen(self, reader, writer):
607606
"""Task that listens to iopub messages."""
608607
try:
609608
_LOGGER.debug("iopub_listen connected")
610-
await self.housekeep_q.put(["register", "iopub", current_task()])
609+
await self.housekeep_q.put(["register", "iopub", asyncio.current_task()])
611610
iopub_socket = ZmqSocket(reader, writer, "PUB")
612611
await iopub_socket.handshake()
613612
self.iopub_socket.add(iopub_socket)
@@ -617,7 +616,7 @@ async def iopub_listen(self, reader, writer):
617616
except asyncio.CancelledError: # pylint: disable=try-except-raise
618617
raise
619618
except EOFError:
620-
await self.housekeep_q.put(["unregister", "iopub", current_task()])
619+
await self.housekeep_q.put(["unregister", "iopub", asyncio.current_task()])
621620
iopub_socket.close()
622621
self.iopub_socket.discard(iopub_socket)
623622
_LOGGER.debug("iopub_listen got eof")
@@ -669,7 +668,7 @@ async def housekeep_run(self):
669668

670669
async def startup_timeout(self):
671670
"""Shut down the session if nothing connects after 30 seconds."""
672-
await self.housekeep_q.put(["register", "startup_timeout", current_task()])
671+
await self.housekeep_q.put(["register", "startup_timeout", asyncio.current_task()])
673672
await asyncio.sleep(30)
674673
if self.task_cnt_max == 1:
675674
#
@@ -680,7 +679,7 @@ async def startup_timeout(self):
680679
self.session_cleanup_callback()
681680
self.session_cleanup_callback = None
682681
await self.housekeep_q.put(["shutdown"])
683-
await self.housekeep_q.put(["unregister", "startup_timeout", current_task()])
682+
await self.housekeep_q.put(["unregister", "startup_timeout", asyncio.current_task()])
684683

685684
async def start_one_server(self, callback):
686685
"""Start a server by finding an available port."""

0 commit comments

Comments
 (0)