def test_docstring_parser2():
example = """
Will block thread until `num_tokens` could be consumed from token bucket `key`.
Args:
key (str): identifying a specific token bucket
num_tokens (int): will block without consuming any tokens until
this amount are availabe to be consumed
timeout (int): seconds to block for
retry_interval (Optional[float]): how long to wait between polling
for tokens to be available. `None` means use default interval
which is equal to time needed to replenish `num_tokens`.
Returns:
Tuple[
int,
str,
ClassName,
]
"""
expected = TypeSignature.factory(
arg_types=ArgTypes.factory(
name=ArgsSection.ARGS,
args=OrderedDict(
[
("key", TypeDef.from_tuples((5, 17), ("str", []), (5, 20),)),
("num_tokens", TypeDef.from_tuples((6, 24), ("int", []), (6, 27),)),
("timeout", TypeDef.from_tuples((8, 21), ("int", []), (8, 24),)),
(
"retry_interval",
TypeDef.from_tuples(
(9, 28), ("Optional", [TypeAtom("float", [])]), (9, 43),
),
),
]
),
),
return_type=ReturnType.factory(
name=ReturnsSection.RETURNS,
type_def=TypeDef.from_tuples(
(15, 12),
(
"Tuple",
[
TypeAtom("int", []),
TypeAtom("str", []),
TypeAtom("ClassName", []),
],
),
(19, 13),
),
),
)
result = docstring_parser.parse(example)
assert result == expected
def test_docstring_parser_invalid_blocks():
example = """
Will block thread until `num_tokens` could be consumed from token bucket `key`.
Args:
this is not an arg: "banana!"
Returns:
send it back within: 7 days
"""
expected = TypeSignature.factory(arg_types=None, return_type=None,)
result = docstring_parser.parse(example)
assert result == expected
def test_remove_types(docstring):
example, context = docstring
signature = docstring_parser.parse(example)
result = remove_types(example, signature)
expected = context["example_no_type"]
note(f"signature: {pprint.pformat(signature)}")
note(f"context: {pprint.pformat(context)}")
note(f"example: {example!r}")
note(f"result: {result!r}")
note(f"expected: {expected!r}")
assert result == expected
def test_docstring_parser_no_annotations():
example = """
Will block thread until `num_tokens` could be consumed from token bucket `key`.
key (str): identifying a specific token bucket
bool: whether we got the requested tokens or not
(False if timed out)
This is not Napoleon format.
"""
expected = TypeSignature.factory(arg_types=None, return_type=None,)
result = docstring_parser.parse(example)
assert result == expected
def test_docstring_parser(docstring):
example, context = docstring
result = docstring_parser.parse(example)
if context["args_section"][0]:
note(f"args_section: {context['args_section'][0]}")
_validate_args_section(
example, result.arg_types, context["args_section"].context,
)
if context["returns_section"][0]:
note(f"returns_section: {context['returns_section'][0]}")
note(result)
_validate_returns_section(
example, result.return_type, context["returns_section"].context,
)
def test_docstring_parser():
example = """
Will block thread until `num_tokens` could be consumed from token bucket `key`.
Keyword Arguments:
key (str): identifying a specific token bucket
num_tokens (int): will block without consuming any tokens until
this amount are availabe to be consumed
timeout (int): seconds to block for
retry_interval (Optional[float]): how long to wait between polling
for tokens to be available. `None` means use default interval
which is equal to time needed to replenish `num_tokens`.
**kwargs (Any)
Returns:
bool: whether we got the requested tokens or not
(False if timed out)
"""
expected = TypeSignature.factory(
arg_types=ArgTypes.factory(
name=ArgsSection.ARGS,
args=OrderedDict(
[
("key", TypeDef.from_tuples((4, 17), ("str", []), (4, 20),)),
("num_tokens", TypeDef.from_tuples((5, 24), ("int", []), (5, 27),)),
("timeout", TypeDef.from_tuples((7, 21), ("int", []), (7, 24),)),
(
"retry_interval",
TypeDef.from_tuples(
(8, 28), ("Optional", [TypeAtom("float", [])]), (8, 43),
),
),
("**kwargs", TypeDef.from_tuples((11, 22), ("Any", []), (11, 25),)),
]
),
),
return_type=ReturnType.factory(
name=ReturnsSection.RETURNS,
type_def=TypeDef.from_tuples((14, 12), ("bool", []), (14, 16),),
),
)
result = docstring_parser.parse(example)
assert result == expected
def m_add_type_comment(node: LN, capture: Capture, filename: Filename,
threadlocals) -> LN:
"""
(modifier)
Adds type comment annotations for functions, as understood by
`mypy --py2` type checking.
"""
threadlocals.docstring_count += 1
# since we filtered for funcs with a docstring, the initial_indent_node
# should be the indent before the start of the docstring quotes.
initial_indent = capture["initial_indent_node"]
function: Leaf = capture["function_name"]
try:
doc_annotation = docstring_parser.parse(
capture["docstring_node"].value)
except parsy.ParseError as e:
report_parse_error(e, function)
raise Interrupt
if not doc_annotation.has_types:
raise Interrupt
annotation_arg_names = (doc_annotation.arg_types.args.keys()
if doc_annotation.arg_types else set())
signature_name, signature_arg_names = threadlocals.signatures[
node.get_lineno()]
assert signature_name == function.value, f"{signature_name} != {function.value}"
if doc_annotation.arg_types and not annotation_arg_names == set(
signature_arg_names):
report_doc_args_signature_mismatch_error(function)
raise Interrupt
# we either have no annotation args, or we do and the names match the signature
# are we okay to annotate?
# TODO these are currently WARN/FAIL... maybe should be OK/WARN/FAIL
# configurably, like for amibiguous types
if signature_arg_names and (not doc_annotation.arg_types or
not doc_annotation.arg_types.is_fully_typed):
report_incomplete_arg_types(function)
if not threadlocals.settings.ALLOW_UNTYPED_ARGS:
raise Interrupt
elif not signature_arg_names and not doc_annotation.arg_types:
# special case: replace doc_annotation with one having empty args
# (rather than `None`)
doc_annotation = TypeSignature.factory(
arg_types=ArgTypes.no_args_factory(),
return_type=doc_annotation.return_type,
)
if not doc_annotation.return_type or not doc_annotation.return_type.is_fully_typed:
report_incomplete_return_type(function)
if threadlocals.settings.REQUIRE_RETURN_TYPE:
raise Interrupt
# yes, annotate...
threadlocals.typed_docstring_count += 1
# print(doc_annotation.return_type, doc_annotation.return_type.is_fully_typed, doc_annotation.return_type.name is ReturnsSection.YIELDS)
if (doc_annotation.return_type
and doc_annotation.return_type.is_fully_typed
and doc_annotation.return_type.name is ReturnsSection.YIELDS):
report_generator_annotation(function)
# record the types we found in this docstring
# and warn/fail on ambiguous types according to IMPORT_COLLISION_POLICY
name_to_strategy: Dict[str, ImportStrategy] = {}
for name in doc_annotation.type_names():
try:
name_to_strategy[
name] = threadlocals.import_strategist.get_for_name(name)
except AmbiguousTypeError as e:
report_ambiguous_type_error(e, function)
if e.should_fail:
raise Interrupt
record_type_names(name_to_strategy)
# add the type comment as first line of func body (before docstring)
type_comment = get_type_comment(doc_annotation, name_to_strategy)
initial_indent.prefix = f"{initial_indent}{type_comment}\n"
threadlocals.comment_count += 1
# remove types from docstring
new_docstring_node = capture["docstring_node"].clone()
new_docstring_node.value = remove_types(
docstring=capture["docstring_node"].value,
signature=doc_annotation,
)
capture["docstring_node"].replace(new_docstring_node)
return node
def test_remove_types(example, expected):
signature = docstring_parser.parse(example)
result = remove_types(example, signature)
print(repr(result))
assert result == expected