fix: refactor type hint retrieval (#312)

dandhlee · web-flow · commit 0c91f60fb2ea · 2023-10-29T00:34:14.000-04:00
* fix: refactor type hint retrieval

* fix: update spacing

* test: update goldens

* fix: update prefix start/end removal

* fix: update docstring for Annotations
diff --git a/docfx_yaml/extension.py b/docfx_yaml/extension.py
@@ -32,7 +32,7 @@
 from pathlib import Path
 from functools import partial
 from itertools import zip_longest
-from typing import Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Optional
 from black import InvalidInput
 
 try:
@@ -732,6 +732,54 @@ def extract_product_name(name):
     return product_name
 
 
+def _extract_type_name(annotation: Any) -> str:
+    """Extracts the type name for the given inspected object.
+
+    Used to identify and extract the type hints given through inspecting the
+    source code. Carefully extracts only the relevant part for the given
+    annotation.
+
+    Args:
+        annotation: the inspected object in its type format. The type hint used
+            is `Any`, because it's the type of the object inspected itself,
+            which can come as any type available.
+
+    Returns:
+        The extracted type hint in human-readable string format.
+    """
+    type_name = ""
+    # Extract names for simple types.
+    try:
+        return annotation.__name__
+    except AttributeError:
+        pass
+
+    # Try to extract names for more complicated types.
+    type_name = str(annotation)
+    if not annotation.__args__:
+        return type_name
+
+    # If ForwardRef references are found, recursively remove them.
+    prefix_to_remove_start = "ForwardRef('"
+    if prefix_to_remove_start not in type_name:
+        return type_name
+
+    prefix_to_remove_end = "')"
+    prefix_start_len = len(prefix_to_remove_start)
+    prefix_end_len = len(prefix_to_remove_end)
+
+    while prefix_to_remove_start in type_name:
+        start_index = type_name.find(prefix_to_remove_start)
+        end_index = type_name.find(prefix_to_remove_end, start_index)
+        type_name = ''.join([
+            type_name[:start_index],
+            type_name[start_index+prefix_start_len:end_index],
+            type_name[end_index+prefix_end_len:],
+        ])
+
+    return type_name
+
+
 def _create_datam(app, cls, module, name, _type, obj, lines=None):
     """
     Build the data structure for an autodoc class
@@ -767,19 +815,12 @@ def _update_friendly_package_name(path):
                 for annotation in argspec.annotations:
                     if annotation == "return":
                         continue
-                    # Extract names for simple types.
                     try:
-                        type_map[annotation] = (argspec.annotations[annotation]).__name__
-                    # Try to extract names for more complicated types.
+                        type_map[annotation] = _extract_type_name(
+                            argspec.annotations[annotation])
                     except AttributeError:
-                        vartype = argspec.annotations[annotation]
-                        try:
-                            type_map[annotation] = str(vartype._name)
-                            if vartype.__args__:
-                                type_map[annotation] += str(vartype.__args__)[:-2] + ")"
-                        except AttributeError:
-                            print(f"Could not parse argument information for {annotation}.")
-                            continue
+                        print(f"Could not parse argument information for {annotation}.")
+                        continue
 
             # Add up the number of arguments. `argspec.args` contains a list of
             # all the arguments from the function.
diff --git a/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.FeatureSet.yml b/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.FeatureSet.yml
@@ -0,0 +1,19 @@
+### YamlMime:UniversalReference
+api_name: []
+items:
+- children: []
+  class: google.cloud.pubsub_v1.types.FeatureSet
+  fullName: google.cloud.pubsub_v1.types.FeatureSet
+  langs:
+  - python
+  module: google.cloud.pubsub_v1.types
+  name: FeatureSet
+  source:
+    id: FeatureSet
+    path: null
+    startLine: null
+  summary: API documentation for `pubsub_v1.types.FeatureSet` class.
+  syntax: {}
+  type: class
+  uid: google.cloud.pubsub_v1.types.FeatureSet
+references: []
diff --git a/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.FieldOptions.EditionDefault.yml b/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.FieldOptions.EditionDefault.yml
@@ -0,0 +1,19 @@
+### YamlMime:UniversalReference
+api_name: []
+items:
+- children: []
+  class: google.cloud.pubsub_v1.types.FieldOptions.EditionDefault
+  fullName: google.cloud.pubsub_v1.types.FieldOptions.EditionDefault
+  langs:
+  - python
+  module: google.cloud.pubsub_v1.types.FieldOptions
+  name: EditionDefault
+  source:
+    id: EditionDefault
+    path: null
+    startLine: null
+  summary: API documentation for `pubsub_v1.types.FieldOptions.EditionDefault` class.
+  syntax: {}
+  type: class
+  uid: google.cloud.pubsub_v1.types.FieldOptions.EditionDefault
+references: []
diff --git a/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.FieldOptions.yml b/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.FieldOptions.yml
@@ -1,7 +1,8 @@
 ### YamlMime:UniversalReference
 api_name: []
 items:
-- children: []
+- children:
+  - google.cloud.pubsub_v1.types.FieldOptions.EditionDefault
   class: google.cloud.pubsub_v1.types.FieldOptions
   fullName: google.cloud.pubsub_v1.types.FieldOptions
   langs:
@@ -16,4 +17,9 @@ items:
   syntax: {}
   type: class
   uid: google.cloud.pubsub_v1.types.FieldOptions
-references: []
+references:
+- fullName: google.cloud.pubsub_v1.types.FieldOptions.EditionDefault
+  isExternal: false
+  name: EditionDefault
+  parent: google.cloud.pubsub_v1.types.FieldOptions
+  uid: google.cloud.pubsub_v1.types.FieldOptions.EditionDefault
diff --git a/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.yml b/tests/testdata/goldens/gapic-combo/google.cloud.pubsub_v1.types.yml
@@ -28,6 +28,7 @@ items:
   - google.cloud.pubsub_v1.types.EnumValueOptions
   - google.cloud.pubsub_v1.types.ExpirationPolicy
   - google.cloud.pubsub_v1.types.ExtensionRangeOptions
+  - google.cloud.pubsub_v1.types.FeatureSet
   - google.cloud.pubsub_v1.types.FieldDescriptorProto
   - google.cloud.pubsub_v1.types.FieldMask
   - google.cloud.pubsub_v1.types.FieldOptions
@@ -236,6 +237,11 @@ references:
   name: ExtensionRangeOptions
   parent: google.cloud.pubsub_v1.types
   uid: google.cloud.pubsub_v1.types.ExtensionRangeOptions
+- fullName: google.cloud.pubsub_v1.types.FeatureSet
+  isExternal: false
+  name: FeatureSet
+  parent: google.cloud.pubsub_v1.types
+  uid: google.cloud.pubsub_v1.types.FeatureSet
 - fullName: google.cloud.pubsub_v1.types.FieldDescriptorProto
   isExternal: false
   name: FieldDescriptorProto
diff --git a/tests/testdata/goldens/gapic-combo/toc.yml b/tests/testdata/goldens/gapic-combo/toc.yml
@@ -129,11 +129,18 @@
           uid: google.cloud.pubsub_v1.types.ExtensionRangeOptions.Declaration
         name: ExtensionRangeOptions
         uid: google.cloud.pubsub_v1.types.ExtensionRangeOptions
+      - name: FeatureSet
+        uid: google.cloud.pubsub_v1.types.FeatureSet
       - name: FieldDescriptorProto
         uid: google.cloud.pubsub_v1.types.FieldDescriptorProto
       - name: FieldMask
         uid: google.cloud.pubsub_v1.types.FieldMask
-      - name: FieldOptions
+      - items:
+        - name: Overview
+          uid: google.cloud.pubsub_v1.types.FieldOptions
+        - name: EditionDefault
+          uid: google.cloud.pubsub_v1.types.FieldOptions.EditionDefault
+        name: FieldOptions
         uid: google.cloud.pubsub_v1.types.FieldOptions
       - name: FileDescriptorProto
         uid: google.cloud.pubsub_v1.types.FileDescriptorProto
