update: forcing explicit optional types (#17)

src/braket/flake8_plugins/braket_checkstyle_plugin.py

@@ -87,6 +87,7 @@ class _Visitor(ast.NodeVisitor):
         "BCS020": "Function '%s' has return documentation but no return type.",
         "BCS021": "Function '%s' is missing return documentation.",
         "BCS022": "Found '%d' invalid indents starting with line ('%s').",
+        "BCS023": "Argument '%s' defaults to None but type hint doesn't end with '| None'.",
     }
 
     def __init__(self) -> None:
@@ -278,28 +279,49 @@ def _check_annotation(
         node: ast.FunctionDef,
     ) -> None:
         annotation = None
+        default_value = None
         if arg_type == ArgType.DEFAULT:
             if node.args.args[arg_index].annotation:
                 annotation = node.args.args[arg_index].annotation
+            if node.args.defaults:
+                default_index = arg_index - (len(node.args.args) - len(node.args.defaults))
+                if default_index >= 0 and node.args.defaults[default_index]:
+                    default_value = node.args.defaults[default_index]
         elif arg_type == ArgType.KEYWORD:
             if node.args.kwonlyargs[arg_index].annotation:
                 annotation = node.args.kwonlyargs[arg_index].annotation
+            default_index = arg_index - (len(node.args.kwonlyargs) - len(node.args.kw_defaults))
+            if default_index >= 0 and node.args.kw_defaults[default_index]:
+                default_value = node.args.kw_defaults[default_index]
+        documented_type = _remove_all_spaces(arg_hint[1:-1])
         if annotation:
-            documented_type = _remove_all_spaces(arg_hint[1:-1])
             annotation_doc = _remove_all_spaces(self._annotation_to_doc_str(annotation))
             if not _are_type_strings_same(annotation_doc, documented_type):
                 self.add_problem(
                     node=node,
                     code="BCS005",
                     arguments=(arg_name, annotation_doc, documented_type),
                 )
+        if (
+            default_value
+            and isinstance(default_value, ast.Constant)
+            and default_value.value is None
+        ):
+            if not documented_type.endswith("|None") and not documented_type.startswith("Optional"):
+                self.add_problem(
+                    node=node,
+                    code="BCS023",
+                    arguments=(arg_name),
+                )
 
     # flake8: noqa: C901
     def _annotation_to_doc_str(self, annotation) -> str:
         if isinstance(annotation, ast.Name):
             return annotation.id
         if isinstance(annotation, ast.Attribute):
             return annotation.attr
+        if isinstance(annotation, ast.Constant) and annotation.value is None:
+            return "None"
         if isinstance(annotation, ast.Subscript):
             if isinstance(annotation.value, ast.Name):
                 if isinstance(annotation.slice, ast.Index):

test/unit_tests/braket/flake8_plugins/example_files/keyword_functions.py

@@ -20,4 +20,30 @@ def function_1(arg0: str, *, arg1 = 0, arg2: int = 1, arg3: int = 1) -> int:
     Returns:
         int: value of the function.
     """
+    pass
+
+
+def function_2(arg0: str, *, a0: int, a1: str | None = None, a2: str = None, a3: str | None = None) -> int:  # noqa
+    """This is a description.
+    Args:
+        arg0 (str): This is a parameter
+        a0 (int): This is a parameter
+        a1 (str|None): This is an optional parameter
+        a2 (str): This is an optional parameter missing optional type hint
+        a3 (str|None): This is an optional parameter
+    Returns:
+        int: value of the function
+    """
+    pass
+
+
+def function_3(arg0: str, *, a0: int, a1: str) -> int:
+    """This is a description.
+    Args:
+        arg0 (str): This is a parameter
+        a0 (int): This is a keyword parameter
+        a1 (str): This is a keyword parameter
+    Returns:
+        int: value of the function
+    """
     pass

test/unit_tests/braket/flake8_plugins/example_files/missing_doc.py

@@ -1,3 +1,16 @@
 def my_function(a0: float, a1: float) -> float:
     val = a0 + a1 + 0.1
     return val
+
+
+def my_function_2(a0: int, a1: str | None = None, a2: str = None, a3: str | None = None) -> int:
+    """This is a description.
+    Args:
+        a0 (int): This is a parameter
+        a1 (str|None): This is an optional parameter
+        a2 (str): This is an optional parameter missing optional type hint
+        a3 (str|None): This is an optional parameter
+    Returns:
+        int: value of the function
+    """
+    pass

test/unit_tests/braket/flake8_plugins/example_files/more_types.py

@@ -9,12 +9,12 @@ class MyB:
     pass
 
 
-def my_func(a0: int, a1: Callable[[Union[MyA, str]], MyA] = None, a2: Optional[MyB] = None) -> np.ndarray:
+def my_func(a0: int, a1: Callable[[Union[MyA, str]], MyA] | None = None, a2: Optional[MyB] = None) -> np.ndarray:
     """This is a description.
     Args:
         a2 (Optional[MyB]): This is out of order.
-        a1 (Callable[[Union[MyA, str]], MyA]): This is out of order.
+        a1 (Callable[[Union[MyA, str]], MyA] | None): This is out of order.
     Returns:
           This is not indented correctly, and doesn't have the return type.
     """
-    pass
+    pass

test/unit_tests/braket/flake8_plugins/example_files/no_error_functions.py

@@ -161,11 +161,12 @@ def function_15(*args, a0: int) -> int:
     pass
 
 
-def function_16(a0: int | list, a1: bool & float) -> int:
+def function_16(a0: int | list, a1: bool & float, a2: str | None = None) -> int:
     """This is a description.
     Args:
         a0 (int | list): This is a parameter
         a1 (bool & float): This is a parameter also
+        a2 (str | None): This is an optional parameter
     Returns:
         int: value of the function
     """

test/unit_tests/braket/flake8_plugins/test_braket_checkstyle_plugin.py

@@ -38,7 +38,13 @@ def test_no_error_functions() -> None:
             "missing_return_type.py",
             {"1:0 BCS002 - Function 'my_function' is missing a type hint for the return value."},
         ),
-        ("missing_doc.py", {"1:0 BCS003 - Function 'my_function' is missing documentation."}),
+        (
+            "missing_doc.py",
+            {
+                "1:0 BCS003 - Function 'my_function' is missing documentation.",
+                "6:0 BCS023 - Argument 'a2' defaults to None but type hint doesn't end with '| None'.",  # noqa
+            },
+        ),
         (
             "class_functions.py",
             {
@@ -123,6 +129,7 @@ def test_no_error_functions() -> None:
                 "1:0 BCS005 - Argument 'arg2' type hint doesn't match documentation. expected: 'int', documented as: 'bool'.",  # noqa
                 "1:0 BCS011 - Argument 'arg1' is missing type hint documentation.",
                 "1:0 BCS015 - Argument 'arg2' is out of order.",
+                "26:0 BCS023 - Argument 'a2' defaults to None but type hint doesn't end with '| None'.",  # noqa
             },
         ),
     ],