init

.venv/lib/python3.8/site-packages/pyflyby/_comms.py

@@ -0,0 +1,136 @@
+from __future__ import absolute_import, division, print_function
+
+from   pyflyby._log             import logger
+from   pyflyby._imports2s       import SourceToSourceFileImportsTransformation
+from   pyflyby._importstmt      import Import
+import six
+
+# These are comm targets that the frontend (lab/notebook) is expected to
+# open. At this point, we handle only missing imports and
+# formatting imports
+
+MISSING_IMPORTS = "pyflyby.missing_imports"
+FORMATTING_IMPORTS = "pyflyby.format_imports"
+INIT_COMMS = "pyflyby.init_comms"
+PYFLYBY_START_MSG = "# THIS CELL WAS AUTO-GENERATED BY PYFLYBY\n"
+PYFLYBY_END_MSG = "# END AUTO-GENERATED BLOCK\n"
+
+pyflyby_comm_targets= [MISSING_IMPORTS, FORMATTING_IMPORTS]
+
+# A map of the comms opened with a given target name.
+comms = {}
+
+# TODO: Document the expected contract for the different
+# custom comm messages
+
+
+def in_jupyter():
+    from IPython.core.getipython import get_ipython
+    ip = get_ipython()
+    if ip is None:
+        logger.debug("get_ipython() doesn't exist. Comm targets can only"
+                     "be added in an Jupyter notebook/lab/console environment")
+        return False
+    else:
+        try:
+            ip.kernel.comm_manager
+        except AttributeError:
+            logger.debug("Comm targets can only be added in Jupyter "
+                         "notebook/lab/console environment")
+            return False
+        else:
+            return True
+
+
+def _register_target(target_name):
+    from IPython.core.getipython import get_ipython
+    ip = get_ipython()
+    comm_manager = ip.kernel.comm_manager
+    comm_manager.register_target(target_name, comm_open_handler)
+
+
+def initialize_comms():
+    if in_jupyter():
+        for target in pyflyby_comm_targets:
+            _register_target(target)
+        from ipykernel.comm import Comm
+        comm = Comm(target_name=INIT_COMMS)
+        msg = {"type": INIT_COMMS}
+        logger.debug("Requesting frontend to (re-)initialize comms")
+        comm.send(msg)
+
+
+def remove_comms():
+    for target_name, comm in six.iteritems(comms):
+        comm.close()
+        logger.debug("Closing comm for " + target_name)
+
+def send_comm_message(target_name, msg):
+    if in_jupyter():
+        try:
+            comm = comms[target_name]
+        except KeyError:
+            logger.debug("Comm with target_name " + target_name + " hasn't been opened")
+        else:
+            # Help the frontend distinguish between multiple types
+            # of custom comm messages
+            msg["type"] = target_name
+            comm.send(msg)
+            logger.debug("Sending comm message for target " + target_name)
+
+
+def comm_close_handler(comm, message):
+    comm_id = message["comm_id"]
+    for target, comm in six.iterkeys(comms):
+        if comm.comm_id == comm_id:
+            comms.pop(target)
+
+
+def _reformat_helper(input_code, imports):
+    from pyflyby._imports2s import reformat_import_statements
+
+    if PYFLYBY_START_MSG in input_code:
+        before, bmarker, middle = input_code.partition(PYFLYBY_START_MSG)
+    else:
+        before, bmarker, middle = "", "", input_code
+
+    if PYFLYBY_END_MSG in middle:
+        middle, emarker, after = middle.partition(PYFLYBY_END_MSG)
+    else:
+        middle, emarker, after = middle, "", ""
+
+    if imports is not None:
+        transform = SourceToSourceFileImportsTransformation(middle)
+
+        if isinstance(imports, str):
+            imports = [imports]
+
+        for imp in imports:
+            assert isinstance(imp, str)
+            if not imp.strip():
+                continue
+            transform.add_import(Import(imp))
+        middle = str(transform.output())
+
+    return reformat_import_statements(before + bmarker + middle + emarker + after)
+
+def comm_open_handler(comm, message):
+    """
+    Handles comm_open message for pyflyby custom comm messages.
+    https://jupyter-client.readthedocs.io/en/stable/messaging.html#opening-a-comm.
+
+    Handler for all PYFLYBY custom comm messages that are opened by the frontend
+    (at this point, just the jupyterlab frontend does this).
+
+    """
+
+    comm.on_close(comm_close_handler)
+    comms[message["content"]["target_name"]] = comm
+
+    @comm.on_msg
+    def _recv(msg):
+        data = msg["content"]["data"]
+        if data["type"] == FORMATTING_IMPORTS:
+            imports = data.get('imports', None)
+            fmt_code = _reformat_helper(data["input_code"], imports)
+            comm.send({"formatted_code": str(fmt_code), "type": FORMATTING_IMPORTS})