Convert LibCST type Annotation back into type?

67 views Asked by At

After creating a start of a function that converts the Annotation type back into a type:

@typechecked
def extract_type_from_annotation(
    *,
    annotation: Union[
        Subscript, Index, Name, Annotation, Attribute, BinaryOperation
    ],
) -> str:
    """Extract the type information from an annotation.

    Args:
        annotation (Union[Subscript, Index, Name, Annotation]): The annotation
         to extract the type from.

    Returns:
        str: The extracted type information.
    """

    if isinstance(annotation, Subscript):
        # Handle the case where the annotation is a Subscript
        value = extract_type_from_annotation(annotation=annotation.value)
        slice_elements = [
            extract_type_from_annotation(annotation=slice_element.slice)
            for slice_element in annotation.slice
        ]
        return f"{value}[{', '.join(slice_elements)}]"
    if isinstance(annotation, Index):
        something = extract_type_from_annotation(annotation=annotation.value)
        return something
    if isinstance(annotation, Name):
        return str(annotation.value)
    if isinstance(annotation, Annotation):
        something = extract_type_from_annotation(
            annotation=annotation.annotation
        )
        return something
    if isinstance(annotation, Attribute):
        left = extract_type_from_annotation(annotation=annotation.value)
        right = extract_type_from_annotation(annotation=annotation.attr)
        something = f"{left}.{right}"
        return something
    if isinstance(annotation, BinaryOperation):
        left = extract_type_from_annotation(annotation=annotation.left)
        right = extract_type_from_annotation(annotation=annotation.right)
        something = f"{left}.{right}"

        return annotation
    return str(annotation)

I felt I was re-inventing the wheel. I expect this functionality is already built into LibCST, however, I have some difficulties finding it. After looking at these

I had some difficulties applying it to the Annotation object like:

code = Module([]).code_for_node(
                param.Annotation
            )

As that would throw:

libcst._nodes.base.CSTCodegenError: Must specify a concrete default_indicator if default used on indicator.

Question

How to get the type of the parameter Annotation object in LibCST back into a sting of the type?

1

There are 1 answers

0
a.t. On

The answer was to feed the Param object into that snippet, instead of the Annotation object. So:

code = Module([]).code_for_node(
                param
            )

Worked (and yielded: G: nx.DiGraph).

Test

Here is a test that tests the updated function:

"""Tests whether a dummy conversation can be outputted to a file."""
# pylint: disable=R0801
import unittest
from typing import Tuple

import libcst as cst
from libcst import Param, Parameters
from typeguard import typechecked

from jsonmodipy.get.argument_getting import extract_type_from_annotation
from tests.conftest import HardcodedTestdata


class Test_extract_type_from_annotation(unittest.TestCase):
    """Object used to test the get_type_from_param( method."""

    # Initialize test object
    @typechecked
    def __init__(  # type:ignore[no-untyped-def]
        self, *args, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.hardcoded_testdata: HardcodedTestdata = HardcodedTestdata()

    def test_returns_valid_type_for_binary_operator(self) -> None:
        """Verifies the extract_type_from_annotation( function returns a type
        with a binary operator succesfully."""
        function_header: str = """def plot_circular_graph(
    *,
    density: float,
    G: nx.DiGraph,
    recurrent_edge_density: int | float,
    test_scope: Long_scope_of_tests) -> None:
    \"\"\"Hello world.\"\"\""""
        # Parse the function header into libcst.
        source_tree = cst.parse_module(function_header)
        parameters: Parameters = source_tree.body[0].params
        kwonly_params: Tuple[Param, ...] = parameters.kwonly_params
        # Call the function to be tested.
        for param in kwonly_params:
            type_nanotation: str = extract_type_from_annotation(param=param)
            if param.name.value == "density":
                self.assertEqual(type_nanotation, "float")
            if param.name.value == "G":
                self.assertEqual(type_nanotation, "nx.DiGraph")
            if param.name.value == "recurrent_edge_density":
                self.assertEqual(type_nanotation, "int | float")
            if param.name.value == "test_scope":
                self.assertEqual(type_nanotation, "Long_scope_of_tests")

And the full function:

@typechecked
def extract_type_from_annotation(
    *,
    param: Param,
) -> str:
    """Extract the type information from an annotation.

    Args:
        param Param: The parameter to extract the type from.

    Returns:
        str: The extracted type information.
    """
    parameter_text: str = Module([]).code_for_node(
        node=param,
    )

    # Remove any trailing spaces.
    parameter_text = parameter_text.strip()

    # Remove trailing commas if there are any.
    if parameter_text.endswith(","):
        parameter_text = parameter_text[:-1]

    # Split argument name and text into a list and get the type.
    the_type: str = parameter_text.split(":")[1].strip()

    return the_type

To have a complete answer to the question one can merge a filler name with the Annotation object, if one does not have the Param object. This can be done like:

param: Param = Param(name="filler", annotation=param.annotation)