Fix up Record.get()

Author: bryanforbes

asyncpg/mypy/hooks.py

@@ -84,20 +84,16 @@ def record_get(ctx: mypy.plugin.MethodContext) \
                 name = value
 
             if name is None:
+                if default_arg is not None:
+                    return default_arg
+
                 ctx.api.fail('Record "{}" has no key \'{}\''
                              .format(ctx.type.type.name, value),
                              ctx.context)
             else:
                 node = ctx.type.type.get(name)
 
-                assert node is not None
-
-                if node.type is not None:
-                    if default_arg is not None:
-                        return mypy.types.UnionType([node.type, default_arg],
-                                                    ctx.context.line,
-                                                    ctx.context.column)
-                    else:
-                        return node.type
+                if node is not None and node.type is not None:
+                    return node.type
 
     return ctx.default_return_type

asyncpg/protocol/protocol.pyi

@@ -210,9 +210,9 @@ _T = TypeVar('_T')
 
 class Record:
     @overload
-    def get(self, key: Any) -> Optional[Any]: ...
+    def get(self, key: str) -> Optional[Any]: ...
     @overload
-    def get(self, key: Any, default: _T) -> Union[Any, _T]: ...
+    def get(self, key: str, default: _T) -> Union[Any, _T]: ...
     def items(self) -> Iterator[Tuple[str, Any]]: ...
     def keys(self) -> Iterator[str]: ...
     def values(self) -> Iterator[Any]: ...