Add check for legacy default=None optional input

Author: nevillelyh

python/coglet/check.py

@@ -0,0 +1,74 @@
+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:', file=sys.stderr)
+        print('-     x: str=Input(default=None)', file=sys.stderr)
+        print('+     x: Optional[str]=Input(default=None)', file=sys.stderr)
+        print(file=sys.stderr)
+        for l in none_defaults:
+            print(l, file=sys.stderr)
+
+        # Check for `from typing import Optional`
+        imports = find(root.body, ast.ImportFrom, 'module', 'typing')
+        if imports is None or 'Optional' not in [n.name for n in imports.names]:
+            # Missing import, add it at beginning of file or before first import
+            # Skip `#!/usr/bin/env python3` or comments
+            lno = 1
+            while lines[lno - 1].startswith('#'):
+                lno += 1
+            for n in root.body:
+                if type(n) in {ast.Import, ast.ImportFrom}:
+                    lno = n.lineno
+                    break
+            lines = lines[:lno - 1] + ['from typing import Optional'] + lines[lno - 1:]
+    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}'