feat: add snippet index

gapic/samplegen_utils/snippet_index.py

@@ -0,0 +1,175 @@
+# Copyright 2022 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 typing import Optional, Dict
+import re
+
+from google.protobuf import json_format
+
+from gapic.schema import api, metadata
+from gapic.samplegen_utils import snippet_metadata_pb2  # type: ignore
+from gapic.samplegen_utils import types
+
+
+CLIENT_INIT_RE = re.compile(r"^\s+# Create a client")
+REQUEST_INIT_RE = re.compile(r"^\s+# Initialize request argument\(s\)")
+REQUEST_EXEC_RE = re.compile(r"^\s+# Make the request")
+RESPONSE_HANDLING_RE = re.compile(r"^\s+# Handle response")
+
+
+class Snippet:
+    """A single snippet and its metadata.
+
+    Attributes:
+        sample_str (str): The full text of the code snippet.
+        metadata (snippet_metadata_pb2.Snippet): The snippet's metadata.
+    """
+
+    def __init__(self, sample_str: str, sample_metadata):
+        self.sample_str = sample_str
+        self.metadata = sample_metadata
+        self._parse_snippet_segments()
+
+    def _parse_snippet_segments(self):
+        """Parse sections of the sample string and update metadata"""
+        self.sample_lines = self.sample_str.splitlines(keepends=True)
+
+        self._full_snippet = snippet_metadata_pb2.Snippet.Segment(
+            type=snippet_metadata_pb2.Snippet.Segment.SegmentType.FULL)
+        self._short_snippet = snippet_metadata_pb2.Snippet.Segment(
+            type=snippet_metadata_pb2.Snippet.Segment.SegmentType.SHORT)
+        self._client_init = snippet_metadata_pb2.Snippet.Segment(
+            type=snippet_metadata_pb2.Snippet.Segment.SegmentType.CLIENT_INITIALIZATION)
+        self._request_init = snippet_metadata_pb2.Snippet.Segment(
+            type=snippet_metadata_pb2.Snippet.Segment.SegmentType.REQUEST_INITIALIZATION)
+        self._request_exec = snippet_metadata_pb2.Snippet.Segment(
+            type=snippet_metadata_pb2.Snippet.Segment.SegmentType.REQUEST_EXECUTION)
+        self._response_handling = snippet_metadata_pb2.Snippet.Segment(
+            type=snippet_metadata_pb2.Snippet.Segment.SegmentType.RESPONSE_HANDLING,
+            end=len(self.sample_lines)
+        )
+
+        # Index starts at 1 since these represent line numbers
+        for i, line in enumerate(self.sample_lines, start=1):
+            if line.startswith("# [START"):  # do not include region tag lines
+                self._full_snippet.start = i + 1
+                self._short_snippet.start = self._full_snippet.start
+            elif line.startswith("# [END"):
+                self._full_snippet.end = i - 1
+                self._short_snippet.end = self._full_snippet.end
+            elif CLIENT_INIT_RE.match(line):
+                self._client_init.start = i
+            elif REQUEST_INIT_RE.match(line):
+                self._client_init.end = i - 1
+                self._request_init.start = i
+            elif REQUEST_EXEC_RE.match(line):
+                self._request_init.end = i - 1
+                self._request_exec.start = i
+            elif RESPONSE_HANDLING_RE.match(line):
+                self._request_exec.end = i - 1
+                self._response_handling.start = i
+
+        self.metadata.segments.extend([self._full_snippet, self._short_snippet, self._client_init,
+                                      self._request_init, self._request_exec, self._response_handling])
+
+    @property
+    def full_snippet(self) -> str:
+        """The portion between the START and END region tags."""
+        start_idx = self._full_snippet.start - 1
+        end_idx = self._full_snippet.end
+        return "".join(self.sample_lines[start_idx:end_idx])
+
+
+class SnippetIndex:
+    """An index of all the snippets for an API.
+
+    Attributes:
+        metadata_index (snippet_metadata_pb2.Index): The snippet metadata index.
+    """
+
+    def __init__(self, api_schema: api.API):
+        self.metadata_index = snippet_metadata_pb2.Index()  # type: ignore
+
+        # Construct a dictionary to insert samples into based on the API schema
+        # NOTE: In the future we expect the generator to support configured samples,
+        # which will result in more than one sample variant per RPC. At that
+        # time a different data structure (and re-writes of add_snippet and get_snippet)
+        # will be needed.
+        self._index: Dict[str, Dict[str, Dict[str, Optional[Snippet]]]] = {}
+
+        self._index = {
+            s.name: {m: {"sync": None, "async": None} for m in s.methods}
+            for s in api_schema.services.values()
+        }
+
+    def add_snippet(self, snippet: Snippet) -> None:
+        """Add a single snippet to the snippet index.
+
+        Args:
+            snippet (Snippet): The code snippet to be added.
+
+        Raises:
+            UnknownService: If the service indicated by the snippet metadata is not found.
+            RpcMethodNotFound: If the method indicated by the snippet metadata is not found.
+        """
+        service_name = snippet.metadata.client_method.method.service.short_name
+        rpc_name = snippet.metadata.client_method.method.full_name
+
+        service = self._index.get(service_name)
+        if service is None:
+            raise types.UnknownService(
+                "API does not have a service named '{}'.".format(service_name))
+
+        method = service.get(rpc_name)
+        if method is None:
+            raise types.RpcMethodNotFound(
+                "API does not have method '{}' in service '{}'".format(rpc_name, service_name))
+
+        if getattr(snippet.metadata.client_method, "async"):
+            method["async"] = snippet
+        else:
+            method["sync"] = snippet
+
+        self.metadata_index.snippets.append(snippet.metadata)
+
+    def get_snippet(self, service_name: str, rpc_name: str, sync: bool = True) -> Optional[Snippet]:
+        """Fetch a single snippet from the index.
+
+        Args:
+            service_name (str): The name of the service.
+            rpc_name (str): The name of the RPC.
+            sync (bool): True for the sync version of the snippet, False for the async version.
+
+        Returns:
+            Optional[Snippet]: The snippet if it exists, or None.
+
+        Raises:
+            UnknownService: If the service is not found.
+            RpcMethodNotFound: If the method is not found.
+        """
+        # Fetch a snippet from the snippet metadata index
+        service = self._index.get(service_name)
+        if service is None:
+            raise types.UnknownService(
+                "API does not have a service named '{}'.".format(service_name))
+        method = service.get(rpc_name)
+        if method is None:
+            raise types.RpcMethodNotFound(
+                "API does not have method '{}' in service '{}'".format(rpc_name, service_name))
+
+        return method["sync" if sync else "async"]
+
+    def get_metadata_json(self) -> str:
+        """JSON representation of Snippet Index."""
+        return json_format.MessageToJson(self.metadata_index, sort_keys=True)