Use channel get_msg helper method

Author: davidbrochart

ipykernel/inprocess/client.py

@@ -178,6 +178,18 @@ def _dispatch_to_kernel(self, msg):
         idents, reply_msg = self.session.recv(stream, copy=False)
         self.shell_channel.call_handlers_later(reply_msg)
 
+    def get_shell_msg(self, block=True, timeout=None):
+        return  self.shell_channel.get_msg(block, timeout)
+
+    def get_iopub_msg(self, block=True, timeout=None):
+        return  self.iopub_channel.get_msg(block, timeout)
+
+    def get_stdin_msg(self, block=True, timeout=None):
+        return  self.stdin_channel.get_msg(block, timeout)
+
+    def get_control_msg(self, block=True, timeout=None):
+        return  self.control_channel.get_msg(block, timeout)
+
 
 #-----------------------------------------------------------------------------
 # ABC Registration

ipykernel/inprocess/tests/test_kernel.py

@@ -65,7 +65,7 @@ def test_pylab(self):
         """Does %pylab work in the in-process kernel?"""
         kc = self.kc
         kc.execute('%pylab')
-        out, err = assemble_output(kc.iopub_channel)
+        out, err = assemble_output(kc.get_iopub_msg)
         self.assertIn('matplotlib', out)
 
     def test_raw_input(self):
@@ -96,7 +96,7 @@ def test_stdout(self):
         kc = BlockingInProcessKernelClient(kernel=kernel, session=kernel.session)
         kernel.frontends.append(kc)
         kc.execute('print("bar")')
-        out, err = assemble_output(kc.iopub_channel)
+        out, err = assemble_output(kc.get_iopub_msg)
         assert out == 'bar\n'
 
     def test_getpass_stream(self):

ipykernel/tests/test_kernel.py

@@ -29,7 +29,7 @@ def _check_master(kc, expected=True, stream="stdout"):
     execute(kc=kc, code="import sys")
     flush_channels(kc)
     msg_id, content = execute(kc=kc, code="print (sys.%s._is_master_process())" % stream)
-    stdout, stderr = assemble_output(kc.iopub_channel)
+    stdout, stderr = assemble_output(kc.get_iopub_msg)
     assert stdout.strip() == repr(expected)
 
 
@@ -46,7 +46,7 @@ def test_simple_print():
     with kernel() as kc:
         iopub = kc.iopub_channel
         msg_id, content = execute(kc=kc, code="print ('hi')")
-        stdout, stderr = assemble_output(iopub)
+        stdout, stderr = assemble_output(kc.get_iopub_msg)
         assert stdout == 'hi\n'
         assert stderr == ''
         _check_master(kc, expected=True)
@@ -56,7 +56,7 @@ def test_sys_path():
     """test that sys.path doesn't get messed up by default"""
     with kernel() as kc:
         msg_id, content = execute(kc=kc, code="import sys; print(repr(sys.path))")
-        stdout, stderr = assemble_output(kc.iopub_channel)
+        stdout, stderr = assemble_output(kc.get_iopub_msg)
     # for error-output on failure
     sys.stderr.write(stderr)
 
@@ -69,7 +69,7 @@ def test_sys_path_profile_dir():
 
     with new_kernel(['--profile-dir', locate_profile('default')]) as kc:
         msg_id, content = execute(kc=kc, code="import sys; print(repr(sys.path))")
-        stdout, stderr = assemble_output(kc.iopub_channel)
+        stdout, stderr = assemble_output(kc.get_iopub_msg)
     # for error-output on failure
     sys.stderr.write(stderr)
 
@@ -100,7 +100,7 @@ def test_subprocess_print():
         ])
 
         msg_id, content = execute(kc=kc, code=code)
-        stdout, stderr = assemble_output(iopub)
+        stdout, stderr = assemble_output(kc.get_iopub_msg)
         nt.assert_equal(stdout.count("hello"), np, stdout)
         for n in range(np):
             nt.assert_equal(stdout.count(str(n)), 1, stdout)
@@ -124,7 +124,7 @@ def test_subprocess_noprint():
         ])
 
         msg_id, content = execute(kc=kc, code=code)
-        stdout, stderr = assemble_output(iopub)
+        stdout, stderr = assemble_output(kc.get_iopub_msg)
         assert stdout == ''
         assert stderr == ''
 
@@ -150,7 +150,7 @@ def test_subprocess_error():
         ])
 
         msg_id, content = execute(kc=kc, code=code)
-        stdout, stderr = assemble_output(iopub)
+        stdout, stderr = assemble_output(kc.get_iopub_msg)
         assert stdout == ''
         assert "ValueError" in stderr
 
@@ -176,7 +176,7 @@ def test_raw_input():
         kc.input(text)
         reply = kc.get_shell_msg(block=True, timeout=TIMEOUT)
         assert reply['content']['status'] == 'ok'
-        stdout, stderr = assemble_output(iopub)
+        stdout, stderr = assemble_output(kc.get_iopub_msg)
         assert stdout == text + "\n"
 
 
@@ -318,14 +318,14 @@ def test_unc_paths():
         kc.execute("cd {0:s}".format(unc_file_path))
         reply = kc.get_shell_msg(block=True, timeout=TIMEOUT)
         assert reply['content']['status'] == 'ok'
-        out, err = assemble_output(iopub)
+        out, err = assemble_output(kc.get_iopub_msg)
         assert unc_file_path in out
 
         flush_channels(kc)
         kc.execute(code="ls")
         reply = kc.get_shell_msg(block=True, timeout=TIMEOUT)
         assert reply['content']['status'] == 'ok'
-        out, err = assemble_output(iopub)
+        out, err = assemble_output(kc.get_iopub_msg)
         assert 'unc.txt' in out
 
         kc.execute(code="cd")

ipykernel/tests/test_message_spec.py

@@ -288,21 +288,21 @@ def test_execute_silent():
     msg_id, reply = execute(code='x=1', silent=True)
 
     # flush status=idle
-    status = KC.iopub_channel.get_msg(timeout=TIMEOUT)
+    status = KC.get_iopub_msg(timeout=TIMEOUT)
     validate_message(status, 'status', msg_id)
     assert status['content']['execution_state'] == 'idle'
 
-    nt.assert_raises(Empty, KC.iopub_channel.get_msg, timeout=0.1)
+    nt.assert_raises(Empty, KC.get_iopub_msg, timeout=0.1)
     count = reply['execution_count']
 
     msg_id, reply = execute(code='x=2', silent=True)
 
     # flush status=idle
-    status = KC.iopub_channel.get_msg(timeout=TIMEOUT)
+    status = KC.get_iopub_msg(timeout=TIMEOUT)
     validate_message(status, 'status', msg_id)
     assert status['content']['execution_state'] == 'idle'
 
-    nt.assert_raises(Empty, KC.iopub_channel.get_msg, timeout=0.1)
+    nt.assert_raises(Empty, KC.get_iopub_msg, timeout=0.1)
     count_2 = reply['execution_count']
     assert count_2 == count
 
@@ -314,7 +314,7 @@ def test_execute_error():
     assert reply['status'] == 'error'
     assert reply['ename'] == 'ZeroDivisionError'
 
-    error = KC.iopub_channel.get_msg(timeout=TIMEOUT)
+    error = KC.get_iopub_msg(timeout=TIMEOUT)
     validate_message(error, 'error', msg_id)
 
 
@@ -560,7 +560,7 @@ def test_stream():
 
     msg_id, reply = execute("print('hi')")
 
-    stdout = KC.iopub_channel.get_msg(timeout=TIMEOUT)
+    stdout = KC.get_iopub_msg(timeout=TIMEOUT)
     validate_message(stdout, 'stream', msg_id)
     content = stdout['content']
     assert content['text'] == 'hi\n'
@@ -571,7 +571,7 @@ def test_display_data():
 
     msg_id, reply = execute("from IPython.display import display; display(1)")
 
-    display = KC.iopub_channel.get_msg(timeout=TIMEOUT)
+    display = KC.get_iopub_msg(timeout=TIMEOUT)
     validate_message(display, 'display_data', parent=msg_id)
     data = display['content']['data']
     assert data['text/plain'] == '1'

ipykernel/tests/utils.py

@@ -43,10 +43,10 @@ def flush_channels(kc=None):
 
     if kc is None:
         kc = KC
-    for channel in (kc.shell_channel, kc.iopub_channel):
+    for get_msg in (kc.get_shell_msg, kc.get_iopub_msg):
         while True:
             try:
-                msg = channel.get_msg(block=True, timeout=0.1)
+                msg = get_msg(block=True, timeout=0.1)
             except Empty:
                 break
             else:
@@ -149,12 +149,12 @@ def new_kernel(argv=None):
         kwargs['extra_arguments'] = argv
     return manager.run_kernel(**kwargs)
 
-def assemble_output(iopub):
+def assemble_output(get_msg):
     """assemble stdout/err from an execution"""
     stdout = ''
     stderr = ''
     while True:
-        msg = iopub.get_msg(block=True, timeout=1)
+        msg = get_msg(block=True, timeout=1)
         msg_type = msg['msg_type']
         content = msg['content']
         if msg_type == 'status' and content['execution_state'] == 'idle':
@@ -174,7 +174,7 @@ def assemble_output(iopub):
 
 def wait_for_idle(kc):
     while True:
-        msg = kc.iopub_channel.get_msg(block=True, timeout=1)
+        msg = kc.get_iopub_msg(block=True, timeout=1)
         msg_type = msg['msg_type']
         content = msg['content']
         if msg_type == 'status' and content['execution_state'] == 'idle':