Skip to content

Commit

Permalink
More complete type inference for fstrings, etc. (apache#33714)
Browse files Browse the repository at this point in the history
* More complete type inference for fstrings, etc.

This also fixes a bug where literals rather than their types
were loaded onto the stack.

* yapf
  • Loading branch information
robertwb authored and tomstepp committed Feb 3, 2025
1 parent 624e29c commit fef2c54
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 3 deletions.
23 changes: 21 additions & 2 deletions sdks/python/apache_beam/typehints/opcodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,7 @@ def binary_slice(state, args):


def store_slice(state, args):
"""Clears elements off the stack like it was constructing a
"""Clears elements off the stack like it was constructing a
container, but leaves the container type back at stack[-1]
since that's all that is relevant for type checking.
"""
Expand Down Expand Up @@ -344,6 +344,10 @@ def list_to_tuple(state, arg):
state.stack.append(Tuple[element_type(base), ...])


def build_string(state, arg):
state.stack[-arg:] = [str]


def list_extend(state, arg):
tail = state.stack.pop()
base = state.stack[-arg]
Expand Down Expand Up @@ -497,7 +501,7 @@ def load_closure(state, arg):
# See https://docs.python.org/3/library/dis.html#opcode-LOAD_CLOSURE
if (sys.version_info.major, sys.version_info.minor) >= (3, 11):
arg -= len(state.co.co_varnames)
state.stack.append(state.get_closure(arg))
state.stack.append(state.closure_type(arg))


def load_deref(state, arg):
Expand Down Expand Up @@ -554,6 +558,21 @@ def build_slice(state, arg):
state.stack[-arg:] = [slice] # a slice object


def format_value(state, arg):
if arg & 0x04:
state.stack.pop()
state.stack.pop()
state.stack.append(str)


def format_simple(state, arg):
state.stack[-1:][str]


def format_with_spec(state, arg):
state.stack[-2:][str]


def _unpack_lists(state, arg):
"""Extract inner types of Lists and Tuples.
Expand Down
5 changes: 4 additions & 1 deletion sdks/python/apache_beam/typehints/row_type.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,10 @@ def __repr__(self):
'%s=%s' % (name, repr(t)) for name, t in self._fields)

def get_type_for(self, name):
return dict(self._fields)[name]
try:
return dict(self._fields)[name]
except KeyError:
return typehints.Any


class GeneratedClassRowTypeConstraint(RowTypeConstraint):
Expand Down
9 changes: 9 additions & 0 deletions sdks/python/apache_beam/typehints/trivial_inference_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -479,6 +479,15 @@ def testRowAttr(self):
lambda row: (row.x, getattr(row, 'y')),
[row_type.RowTypeConstraint.from_fields([('x', int), ('y', str)])])

def testRowMissingAttr(self):
self.assertReturnType(
typehints.Any,
lambda row: getattr(row, '_asdict'),
[row_type.RowTypeConstraint.from_fields([('x', int), ('y', str)])])

def testFString(self):
self.assertReturnType(str, lambda x, y: f'{x}: {y:0.2}', [str, float])

def testPyCallable(self):
self.assertReturnType(
typehints.Tuple[int, str],
Expand Down

0 comments on commit fef2c54

Please sign in to comment.