feat: implement native asyncio support via Cross-Sync

.cross_sync/README.md

@@ -0,0 +1,75 @@
+# CrossSync
+
+CrossSync provides a simple way to share logic between async and sync code.
+It is made up of a small library that provides:
+1. a set of shims that provide a shared sync/async API surface
+2. annotations that are used to guide generation of a sync version from an async class
+
+Using CrossSync, the async code is treated as the source of truth, and sync code is generated from it.
+
+## Usage
+
+### CrossSync Shims
+
+Many Asyncio components have direct, 1:1 threaded counterparts for use in non-asyncio code. CrossSync
+provides a compatibility layer that works with both
+
+| CrossSync | Asyncio Version | Sync Version |
+| --- | --- | --- |
+| CrossSync.Queue | asyncio.Queue | queue.Queue |
+| CrossSync.Condition | asyncio.Condition | threading.Condition |
+| CrossSync.Future | asyncio.Future | Concurrent.futures.Future |
+| CrossSync.Task | asyncio.Task | Concurrent.futures.Future |
+| CrossSync.Event | asyncio.Event | threading.Event |
+| CrossSync.Semaphore | asyncio.Semaphore | threading.Semaphore |
+| CrossSync.Awaitable | typing.Awaitable | typing.Union (no-op type) |
+| CrossSync.Iterable | typing.AsyncIterable | typing.Iterable |
+| CrossSync.Iterator | typing.AsyncIterator | typing.Iterator |
+| CrossSync.Generator | typing.AsyncGenerator | typing.Generator |
+| CrossSync.Retry | google.api_core.retry.AsyncRetry | google.api_core.retry.Retry |
+| CrossSync.StopIteration | StopAsyncIteration | StopIteration |
+| CrossSync.Mock | unittest.mock.AsyncMock | unittest.mock.Mock |
+
+Custom aliases can be added using `CrossSync.add_mapping(class, name)`
+
+Additionally, CrossSync provides method implementations that work equivalently in async and sync code:
+- `CrossSync.sleep()`
+- `CrossSync.gather_partials()`
+- `CrossSync.wait()`
+- `CrossSync.condition_wait()`
+- `CrossSync,event_wait()`
+- `CrossSync.create_task()`
+- `CrossSync.retry_target()`
+- `CrossSync.retry_target_stream()`
+
+### Annotations
+
+CrossSync provides a set of annotations to mark up async classes, to guide the generation of sync code.
+
+- `@CrossSync.convert_sync`
+  - marks classes for conversion. Unmarked classes will be copied as-is
+  - if add_mapping is included, the async and sync classes can be accessed using a shared CrossSync.X alias
+- `@CrossSync.convert`
+  - marks async functions for conversion. Unmarked methods will be copied as-is
+- `@CrossSync.drop`
+  - marks functions or classes that should not be included in sync output
+- `@CrossSync.pytest`
+  - marks test functions. Test functions automatically have all async keywords stripped (i.e., rm_aio is unneeded)
+- `CrossSync.add_mapping`
+  - manually registers a new CrossSync.X alias, for custom types
+- `CrossSync.rm_aio`
+  - Marks regions of the code that include asyncio keywords that should be stripped during generation
+
+### Code Generation
+
+Generation can be initiated using `nox -s generate_sync`
+from the root of the project. This will find all classes with the `__CROSS_SYNC_OUTPUT__ = "path/to/output"` 
+annotation, and generate a sync version of classes marked with `@CrossSync.convert_sync` at the output path.
+
+There is a unit test at `tests/unit/data/test_sync_up_to_date.py` that verifies that the generated code is up to date
+
+## Architecture
+
+CrossSync is made up of two parts:
+- the runtime shims and annotations live in `/google/cloud/bigtable/_cross_sync`
+- the code generation logic lives in `/.cross_sync/` in the repo root

.cross_sync/generate.py

@@ -0,0 +1,111 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+from typing import Sequence
+import ast
+"""
+Entrypoint for initiating an async -> sync conversion using CrossSync
+
+Finds all python files rooted in a given directory, and uses
+transformers.CrossSyncFileProcessor to handle any files marked with
+__CROSS_SYNC_OUTPUT__
+"""
+
+
+def extract_header_comments(file_path) -> str:
+    """
+    Extract the file header. Header is defined as the top-level
+    comments before any code or imports
+    """
+    header = []
+    with open(file_path, "r", encoding="utf-8-sig") as f:
+        for line in f:
+            if line.startswith("#") or line.strip() == "":
+                header.append(line)
+            else:
+                break
+    header.append("\n# This file is automatically generated by CrossSync. Do not edit manually.\n\n")
+    return "".join(header)
+
+
+class CrossSyncOutputFile:
+
+    def __init__(self, output_path: str, ast_tree, header: str | None = None):
+        self.output_path = output_path
+        self.tree = ast_tree
+        self.header = header or ""
+
+    def render(self, with_formatter=True, save_to_disk: bool = True) -> str:
+        """
+        Render the file to a string, and optionally save to disk
+
+        Args:
+            with_formatter: whether to run the output through black before returning
+            save_to_disk: whether to write the output to the file path
+        """
+        full_str = self.header + ast.unparse(self.tree)
+        if with_formatter:
+            import black  # type: ignore
+            import autoflake  # type: ignore
+
+            full_str = black.format_str(
+                autoflake.fix_code(full_str, remove_all_unused_imports=True),
+                mode=black.FileMode(),
+            )
+        if save_to_disk:
+            import os
+            os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
+            with open(self.output_path, "w") as f:
+                f.write(full_str)
+        return full_str
+
+
+def convert_files_in_dir(directory: str) -> set[CrossSyncOutputFile]:
+    import glob
+    from transformers import CrossSyncFileProcessor
+
+    # find all python files in the directory
+    files = glob.glob(directory + "/**/*.py", recursive=True)
+    # keep track of the output files pointed to by the annotated classes
+    artifacts: set[CrossSyncOutputFile] = set()
+    file_transformer = CrossSyncFileProcessor()
+    # run each file through ast transformation to find all annotated classes
+    for file_path in files:
+        ast_tree = ast.parse(open(file_path, encoding="utf-8-sig").read())
+        output_path = file_transformer.get_output_path(ast_tree)
+        if output_path is not None:
+            # contains __CROSS_SYNC_OUTPUT__ annotation
+            converted_tree = file_transformer.visit(ast_tree)
+            header = extract_header_comments(file_path)
+            artifacts.add(CrossSyncOutputFile(output_path, converted_tree, header))
+    # return set of output artifacts
+    return artifacts
+
+
+def save_artifacts(artifacts: Sequence[CrossSyncOutputFile]):
+    for a in artifacts:
+        a.render(save_to_disk=True)
+
+
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) < 2:
+        print("Usage: python .cross_sync/generate.py <directory>")
+        sys.exit(1)
+
+    search_root = sys.argv[1]
+    outputs = convert_files_in_dir(search_root)
+    print(f"Generated {len(outputs)} artifacts: {[a.output_path for a in outputs]}")
+    save_artifacts(outputs)