pyflyby 1.10.4__cp311-cp311-macosx_11_0_arm64.whl

pyflyby/_comms.py

@@ -0,0 +1,221 @@
+
+from __future__ import print_function
+
+from   pyflyby._imports2s       import (SourceToSourceFileImportsTransformation,
+                                        SourceToSourceImportBlockTransformation,
+                                        fix_unused_and_missing_imports,
+                                        replace_star_imports,
+                                        reformat_import_statements)
+from   pyflyby._importstmt      import Import
+from   pyflyby._log             import logger
+
+from   typing                   import Dict, Any
+
+# 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"
+TIDY_IMPORTS = "pyflyby.tidy_imports"
+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, TIDY_IMPORTS]
+
+# A map of the comms opened with a given target name.
+comms:Dict[str, Any] = {}
+
+# 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 comms.items():
+        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 comms.keys():
+        if comm.comm_id == comm_id:
+            comms.pop(target)
+
+
+def _reformat_helper(input_code, imports):
+    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 extract_import_statements(text):
+    """This is a util for notebook interactions and extracts import statements
+    from some python code. This function also re-orders imports.
+    Args:
+        code (str): The code from which import statements have to be extracted
+
+    Returns:
+        (str, str): The first returned value contains all the import statements.
+        The second returned value is the remaining code after
+        extracting the import statements.
+    """
+    transformer = SourceToSourceFileImportsTransformation(text)
+    imports = '\n'.join([str(im.pretty_print()).strip() for im in transformer.import_blocks])
+    remaining_code = "\n".join([str(st.pretty_print()).strip() if not isinstance(st, SourceToSourceImportBlockTransformation) else "" for st in transformer.blocks])
+    return imports, remaining_code
+
+def collect_code_with_imports_on_top(imports: str, cell_array):
+    return (
+        imports
+        + "\n"
+        + "\n".join(
+            [
+                cell["text"] if cell["type"] == "code" and not cell.get("ignore", False) else ""
+                for cell in cell_array
+            ]
+        )
+    )
+
+def run_tidy_imports(code):
+    return str(
+        reformat_import_statements(
+            fix_unused_and_missing_imports(
+                replace_star_imports(code)
+            )
+        )
+    )
+
+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:
+            msg_id = data.get("msg_id", None)
+            imports = data.get("imports", None)
+            fmt_code = _reformat_helper(data["input_code"], imports)
+            comm.send(
+                {
+                    "msg_id": msg_id,
+                    "formatted_code": str(fmt_code),
+                    "type": FORMATTING_IMPORTS,
+                }
+            )
+        elif data["type"] == TIDY_IMPORTS:
+            checksum = data.get("checksum", '')
+            cell_array = data.get("cellArray", [])
+            # import_statements is a string because when
+            # SourceToSourceFileImportsTransformation is run on a piece of code
+            # it will club similar imports together and re-order the imports
+            # by making the imports a string, all the imports are processed
+            # together making sure tidy-imports has context on all the imports
+            # while clubbing similar imports and re-ordering them.
+            import_statements, processed_cell_array = "", []
+            for cell in cell_array:
+                ignore = False
+                text = cell.get("text")
+                cell_type = cell.get("type")
+                if cell_type == "code":
+                    try:
+                        imports, text = extract_import_statements(text)
+                        import_statements += (imports + "\n")
+                    except SyntaxError:
+                        # If a cell triggers Syntax Error, we set ignore to
+                        # True and don't include it when running tidy-imports
+                        # For eg. this is triggered due to cells with magic
+                        # commands
+                        ignore = True
+                processed_cell_array.append({"text": text, "type": cell_type, "ignore": ignore})
+            code_with_collected_imports = collect_code_with_imports_on_top(import_statements, processed_cell_array)
+            code_post_tidy_imports = run_tidy_imports(code_with_collected_imports)
+            import_statements, _ = extract_import_statements(code_post_tidy_imports)
+            comm.send(
+                {
+                    "checksum": checksum,
+                    "type": TIDY_IMPORTS,
+                    "cells": processed_cell_array,
+                    "imports": import_statements,
+                }
+            )