Add check for legacy default=None optional input

Author: nevillelyh

python/coglet/check.py

@@ -0,0 +1,59 @@
+import ast
+import sys
+from typing import Optional, Type, TypeVar
+
+T = TypeVar('T', covariant=True)
+
+
+def find(nodes: list, tpe: Type[T], attr: str, name: str) -> Optional[T]:
+    for n in nodes:
+        if type(n) is tpe and getattr(n, attr) == name:
+            return n
+    return None
+
+
+def check(file: str, predictor: str) -> None:
+    with open(file, 'r') as f:
+        content = f.read()
+    lines = content.splitlines()
+    root = ast.parse(content)
+
+    p = find(root.body, ast.ClassDef, 'name', predictor)
+    if p is None:
+        return
+    fn = find(p.body, ast.FunctionDef, 'name', 'predict')
+    if fn is None:
+        fn = find(p.body, ast.AsyncFunctionDef, 'name', 'predict')  # type: ignore
+    args_and_defaults = zip(fn.args.args[-len(fn.args.defaults) :], fn.args.defaults)  # type: ignore
+    none_defaults = []
+    for a, d in args_and_defaults:
+        if type(a.annotation) is not ast.Name:
+            continue
+        if type(d) is not ast.Call or d.func.id != 'Input':  # type: ignore
+            continue
+        v = find(d.keywords, ast.keyword, 'arg', 'default')
+        if v is None or type(v.value) is not ast.Constant:
+            continue
+        if v.value.value is None:
+            pos = f'{file}:{a.lineno}:{a.col_offset}'
+
+            # Add Optional[] to type annotation
+            # No need to remove `default=None` since `x: Optional[T] = Input(default=None)` is valid
+            ta = a.annotation
+            l = lines[ta.lineno - 1]
+            parts = l[:ta.col_offset], l[ta.col_offset:ta.end_col_offset], l[ta.end_col_offset:]
+            l = f'{parts[0]}Optional[{parts[1]}]{parts[2]}'
+            lines[ta.lineno - 1] = l
+
+            none_defaults.append(f'{pos}: {a.arg}: {ta.id}={ast.unparse(d)}')
+    if len(none_defaults) > 0:
+        print('Default value of None without explicit Optional[T] type hint is ambiguous and deprecated, for example:')
+        print('-     x: str=Input(default=None)')
+        print('+     x: Optional[str]=Input(default=None)')
+        print()
+        for l in none_defaults:
+            print(l)
+    print('\n'.join(lines))
+
+
+check(sys.argv[1], sys.argv[2])

python/tests/cases/legacy_optional.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def predict(
+        self,
+        s1: str,
+        s2: Optional[str],
+        s3: str = Input(),
+        s4: Optional[str] = Input(),
+        s5: str = Input(default=None),
+        s6: str = Input(default=None, description='s6'),
+        s7: str = Input(description='s7', default=None),
+        s8: Optional[str] = Input(default=None),
+    ) -> str:
+        return f'{s1}:{s2}:{s3}:{s4}:{s5}:{s6}'