Any way to get type hints to work for both 2.7/3.6 in VS code?

I'd like to use type hints in a few places. This syntax works with VS code, but isn't compatible with 2.7:

def foo(a: str, b: int) -> str:
    pass

From what I can tell, the following is the suggested way of using hints, but VS code doesn't recognize it:

def foo(a, b):
    # type: (str, int) -> str
    pass

The following works, though i'd hate to have to do it this way:

def foo(a, b):
    a = a  # type: str
    b = b  # type: int
    pass

Is there a better way to do this?