Repository: facebookincubator/Bowler

Issue:
Would be great to have a helper that infers numeric types
For a recent project I wished there were some functions that would let me figure out int vs float (vs unknown) for simple expressions.  I'm not aiming for this to be a full type-checker, or even go beyond those three (int/float/unknown), but I'll run the experiment to see if this ends up being too many lines of code.

Discussion:

Patch:
diff --git a/bowler/helpers.py b/bowler/helpers.py
index 4df80ab..f7fd79a 100644
--- a/bowler/helpers.py
+++ b/bowler/helpers.py
@@ -12,7 +12,7 @@ import click
 from fissix.pytree import Leaf, Node, type_repr
 from fissix.pgen2.token import tok_name
 
-from .types import LN, SYMBOL, Capture, Filename, FilenameMatcher
+from .types import LN, SYMBOL, TOKEN, Capture, Filename, FilenameMatcher
 
 log = logging.getLogger(__name__)
 
@@ -122,6 +122,15 @@ def is_method(node: LN) -> bool:
     )
 
 
+def is_call_to(node: LN, func_name: str) -> bool:
+    """Returns whether the node represents a call to the named function."""
+    return (
+        node.type == SYMBOL.power
+        and node.children[0].type == TOKEN.NAME
+        and node.children[0].value == func_name
+    )
+
+
 def find_first(node: LN, target: int, recursive: bool = False) -> Optional[LN]:
     queue: List[LN] = [node]
     queue.extend(node.children)
diff --git a/bowler/type_inference.py b/bowler/type_inference.py
new file mode 100644
index 0000000..e88dd73
--- /dev/null
+++ b/bowler/type_inference.py
@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""bowler.type_inference
+
+Given an expression, find its result type.
+
+For sufficiently obvious expressions, we can find this using only local
+knowledge (numeric literals, and functions/names which have a standard
+meaning).  Some obvious examples:
+
+`1.0` -> InferredType.FLOAT
+`2L` -> InferredType.INT
+`1/2` -> Depends on use_py2_division
+
+Even in cases where we don't know the full inputs, we can make some reasonable
+assumptions.  For example when passing `use_py2_division=False,
+type_for_unknown=InferredType.INT_OR_FLOAT`:
+
+`x+1.0` -> InferredType.FLOAT
+`x/y` -> InferredType.FLOAT
+`len(x) + 1` -> InferredType.INT
+`float(z)` -> InferredType.FLOAT
+
+This is intended to be useful for either refactoring or flagging for humans
+syntax like `range(float(...))` or `"%d" % (float(...),)`.
+"""
+
+import enum
+from typing import Dict, Sequence, Union
+
+from fissix import pygram, pytree
+from fissix.pgen2 import token
+from fissix.pgen2.driver import Driver
+
+from .types import LN, SYMBOL, TOKEN
+from .helpers import is_call_to
+
+__all__ = ["InferredType", "numeric_expr_type"]
+
+
+class InferredType(enum.IntEnum):
+    # The order of these is important, such that an expression using py3
+    # semantics most operators take max(OP_MIN_TYPE[op], max(children)) as the
+    # result.
+    UNSET = 0
+    BOOL = 1
+    INT = 2
+    # This represents UNKNOWN but assumed to be restricted to normal numeric
+    # values.  It can still be promoted to COMPLEX or FLOAT, but if is the
+    # final result should be treated as INT (or better).
+    INT_OR_FLOAT = 3
+    FLOAT = 4
+    COMPLEX = 5
+    UNKNOWN = 6
+
+
+# Note: SLASH and DOUBLESLASH are specialcased.
+OP_MIN_TYPE: Dict = {
+    TOKEN.PLUS: InferredType.INT,
+    TOKEN.MINUS: InferredType.INT,
+    TOKEN.STAR: InferredType.INT,
+    TOKEN.PERCENT: InferredType.INT,
+    TOKEN.SLASH: InferredType.INT,
+    TOKEN.DOUBLESLASH: InferredType.INT,
+    TOKEN.TILDE: InferredType.INT,  # bitwise not
+    TOKEN.DOUBLESTAR: InferredType.INT,
+    TOKEN.LEFTSHIFT: InferredType.INT,
+    TOKEN.RIGHTSHIFT: InferredType.INT,
+    TOKEN.VBAR: InferredType.BOOL,
+    TOKEN.CIRCUMFLEX: InferredType.BOOL,
+    TOKEN.AMPER: InferredType.BOOL,
+    TOKEN.LESS: InferredType.BOOL,
+}
+
+
+def numeric_expr_type(
+    node: LN,
+    use_py2_division=False,
+    type_for_unknown: InferredType = InferredType.UNKNOWN,
+) -> "InferredType":
+    """Infer the type of an expression from its literals.
+
+    We broaden the definition of "literal" a bit to also include calls to
+    certain functions like int() and float() where the return type does not
+    change based on the arguments.
+
+    Args:
+        node: A Node or leaf.
+        use_py2_division: Whether to use magical python 2 style division.
+        type_for_unknown: An InferredType to customize how you wan unknown
+            handled.  Use `INT_OR_FLOAT` if you trust your input to only work
+            on numbers, but `UNKNOWN` if you want objects to be an option.
+
+    Returns: InferredType
+    """
+    if node.type == TOKEN.NUMBER:
+        # It's important that we do not use eval here; some literals like `2L`
+        # may be invalid in the current interpreter.
+        if "j" in node.value:
+            return InferredType.COMPLEX
+        elif "." in node.value or "e" in node.value:
+            return InferredType.FLOAT
+        return InferredType.INT
+    elif node.type == TOKEN.NAME and node.value in ("True", "False"):
+        return InferredType.BOOL
+    # TODO let the caller provide other known return types, or even a
+    # collection of locals and their types.
+    elif is_call_to(node, "bool"):
+        return InferredType.BOOL
+    elif is_call_to(node, "int") or is_call_to(node, "len"):
+        return InferredType.INT
+    elif is_call_to(node, "float"):
+        return InferredType.FLOAT
+
+    elif node.type in (SYMBOL.comparison, SYMBOL.not_test):
+        return InferredType.BOOL
+    elif node.type == SYMBOL.factor:
+        # unary ~ + -, always [op, number]
+        return max(
+            OP_MIN_TYPE[node.children[0].type],
+            numeric_expr_type(node.children[1], use_py2_division, type_for_unknown),
+        )
+    elif node.type == SYMBOL.shift_expr:
+        # << only valid on int
+        return InferredType.INT
+
+    elif node.type == SYMBOL.power:
+        # a**b, but also f(...)
+        if node.children[1].type != TOKEN.DOUBLESTAR:
+            # probably f(...)
+            return type_for_unknown
+
+        return max(
+            max(OP_MIN_TYPE[c.type] for c in node.children[1::2]),
+            max(
+                numeric_expr_type(c, use_py2_division, type_for_unknown)
+                for c in node.children[::2]
+            ),
+        )
+    elif node.type in (
+        SYMBOL.arith_expr,
+        SYMBOL.xor_expr,
+        SYMBOL.and_expr,
+        SYMBOL.expr,
+    ):
+        return max(
+            max(OP_MIN_TYPE[c.type] for c in node.children[1::2]),
+            max(
+                numeric_expr_type(c, use_py2_division, type_for_unknown)
+                for c in node.children[::2]
+            ),
+        )
+    elif node.type == SYMBOL.term:
+        # */%
+        # This is where things get interesting, as we handle use_py2_division.
+        t = InferredType.UNSET
+        last_op = None
+        for i in range(len(node.children)):
+            if i % 2 == 0:
+                new = numeric_expr_type(
+                    node.children[i], use_py2_division, type_for_unknown
+                )
+                if last_op == TOKEN.DOUBLESLASH:
+                    t = InferredType.INT
+                elif last_op == TOKEN.SLASH:
+                    if use_py2_division:
+                        if t == InferredType.INT and new == InferredType.INT:
+                            t = InferredType.INT
+                        else:
+                            t = max(t, max(OP_MIN_TYPE[last_op], new))
+                    else:
+                        t = max(t, InferredType.FLOAT)
+                else:
+                    if last_op:
+                        t = max(t, OP_MIN_TYPE[last_op])
+                    t = max(t, new)
+            else:
+                last_op = node.children[i].type
+        return t
+    elif node.type in (SYMBOL.or_test, SYMBOL.and_test):
+        return max(
+            numeric_expr_type(c, use_py2_division, type_for_unknown)
+            for c in node.children[::2]
+        )
+
+    return type_for_unknown
diff --git a/makefile b/makefile
index 5b7b7ff..6c1b717 100644
--- a/makefile
+++ b/makefile
@@ -19,6 +19,13 @@ format:
 	black bowler setup.py
 
 lint:
+	@/bin/bash -c 'die() { echo "$$1"; exit 1; }; \
+	  while read filename; do \
+	  grep -q "Copyright (c) Facebook" "$$filename" || \
+	    die "Missing copyright in $$filename"; \
+	  grep -q "#!/usr/bin/env python3" "$$filename" || \
+	    die "Missing #! in $$filename"; \
+	  done < <( git ls-tree -r --name-only HEAD | grep ".py$$" )'
 	black --check bowler setup.py
 	mypy --ignore-missing-imports --python-version 3.6 .
 
diff --git a/setup.py b/setup.py
index d8dfa26..ceb7631 100644
--- a/setup.py
+++ b/setup.py
@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+#
 # Copyright (c) Facebook, Inc. and its affiliates.
 #
 # This source code is licensed under the MIT license found in the


Test patch:
diff --git a/bowler/tests/__init__.py b/bowler/tests/__init__.py
index 25ee3c8..23f9dca 100644
--- a/bowler/tests/__init__.py
+++ b/bowler/tests/__init__.py
@@ -8,3 +8,4 @@
 from .query import QueryTest
 from .smoke import SmokeTest
 from .tool import ToolTest
+from .type_inference import OpMinTypeTest, ExpressionTest
diff --git a/bowler/tests/type_inference.py b/bowler/tests/type_inference.py
new file mode 100644
index 0000000..81edba4
--- /dev/null
+++ b/bowler/tests/type_inference.py
@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import sys
+import unittest
+from io import StringIO
+
+from fissix import pygram, pytree
+from fissix.pgen2.driver import Driver
+
+from ..type_inference import OP_MIN_TYPE, InferredType, numeric_expr_type
+from ..types import LN, SYMBOL, TOKEN
+
+BINARY_OPERATORS = ["+", "-", "*", "**", "<<", ">>", "|", "&", "^", "%", "<"]
+
+UNARY_OPERATORS = ["~", "-", "+"]
+# TODO 'is', 'not'
+
+SAMPLE_EXPRESSIONS = [
+    # and/or
+    ("True", InferredType.BOOL),
+    ("True or False", InferredType.BOOL),
+    ("True or 1", InferredType.INT),
+    ("1 or 1.0", InferredType.FLOAT),
+    # Calls
+    ("bool(x)", InferredType.BOOL),
+    ("int(x)", InferredType.INT),
+    ("len(x)", InferredType.INT),
+    ("float(x)", InferredType.FLOAT),
+    # Basic
+    ("func()", InferredType.INT_OR_FLOAT),
+    ("x+1", InferredType.INT_OR_FLOAT),
+    ("x*1", InferredType.INT_OR_FLOAT),
+    ("1+x", InferredType.INT_OR_FLOAT),
+    ("1*x", InferredType.INT_OR_FLOAT),
+    # Single
+    ("1+1", InferredType.INT),
+    ("1.0+1.0", InferredType.FLOAT),
+    # Mixed
+    ("1+1.0", InferredType.FLOAT),
+    ("1+x/2", InferredType.FLOAT),  # py3
+    ("1.0+x/2", InferredType.FLOAT),
+    # Division
+    ("1/1", InferredType.FLOAT),
+    ("1/2.0", InferredType.FLOAT),
+    ("1.0/2", InferredType.FLOAT),
+    ("1/x", InferredType.FLOAT),
+    ("1.0/x", InferredType.FLOAT),
+    ("True/True", InferredType.FLOAT),
+    ("1j/2", InferredType.COMPLEX),
+    # Floor Division
+    ("1//2", InferredType.INT),
+    ("1.0//2.0", InferredType.INT),
+]
+
+SAMPLE_PY2_EXPRESSIONS = [
+    # Py2 Division
+    ("1/1", InferredType.INT),
+    ("1/2.0", InferredType.FLOAT),
+    ("1.0/2", InferredType.FLOAT),
+    ("1/x", InferredType.INT_OR_FLOAT),
+    ("1.0/x", InferredType.FLOAT),
+    ("True/True", InferredType.INT),
+    ("1j/2", InferredType.COMPLEX),
+]
+
+
+def _produce_test(lcals, gen_func, args):
+    t = gen_func(*args)
+    t.__name__ += str(args)
+    lcals[t.__name__] = t
+
+
+def tree(input: str) -> LN:
+    print(f"Input is {repr(input)}")
+    driver = Driver(pygram.python_grammar_no_print_statement, convert=pytree.convert)
+    return driver.parse_string(input)
+
+
+def map_type(o):
+    if isinstance(o, complex):
+        return InferredType.COMPLEX
+    elif isinstance(o, float):
+        return InferredType.FLOAT
+    elif isinstance(o, bool):
+        return InferredType.BOOL
+    elif isinstance(o, int):
+        return InferredType.INT
+
+
+class OpMinTypeTest(unittest.TestCase):
+    """
+    Verifies that the generated OP_MIN_TYPE matches that of the current Python
+    interpreter, and that `numeric_expr_type` agrees.
+    """
+
+    def _run_test(self, expression_str, expected_type):
+        t = tree(expression_str)
+        expr = t.children[0].children[0]
+        # t_op = expr.children[1].type
+        # expr_type = pytree.type_repr(expr.type)
+        # key_type = TOKEN.tok_name[t_op]
+
+        inferred_result = numeric_expr_type(expr)
+
+        self.assertEqual(expected_type, inferred_result, repr(expr))
+
+    def gen_test_binop(op):
+        def test_binop(self):
+            for lhs in ("True", "1", "1.0", "2j"):
+                for rhs in ("True", "1", "1.0", "2j"):
+                    snippet = f"{lhs} {op} {rhs}\n"
+                    try:
+                        real_result = eval(snippet, {}, {})
+                    except TypeError:
+                        continue
+                    self._run_test(snippet, map_type(real_result))
+
+        return test_binop
+
+    def gen_test_uniop(op):
+        def test_uniop(self):
+            for rhs in ("True", "1", "1.0", "2j"):
+                snippet = f"{op} {rhs}\n"
+                try:
+                    real_result = eval(snippet, {}, {})
+                except TypeError:
+                    continue
+                self._run_test(snippet, map_type(real_result))
+
+        return test_uniop
+
+    def gen_test_min_type(op):
+        def test_min_type(self):
+            snippet = f"True {op} True\n"
+            real_result = eval(snippet, {}, {})
+
+            t = tree(snippet)
+            expr = t.children[0].children[0]
+            t_op = expr.children[1].type
+            expr_type = pytree.type_repr(expr.type)
+            key_type = TOKEN.tok_name[t_op]
+
+            self.assertEqual(map_type(real_result), OP_MIN_TYPE.get(t_op), key_type)
+
+        return test_min_type
+
+    def gen_test_min_type_unary(op):
+        def test_min_type_unary(self):
+            snippet = f"{op} True\n"
+            real_result = eval(snippet, {}, {})
+
+            t = tree(snippet)
+
+            expr = t.children[0].children[0]
+            t_op = expr.children[0].type
+            expr_type = pytree.type_repr(expr.type)
+            key_type = TOKEN.tok_name[t_op]
+
+            self.assertEqual(map_type(real_result), OP_MIN_TYPE.get(t_op), expr_type)
+
+        return test_min_type_unary
+
+    # This produces real methods that can have normal decorators on them, with
+    # a proper pass/fail count (unlike subTest).
+    for i, op in enumerate(BINARY_OPERATORS):
+        _produce_test(locals(), gen_test_binop, (op,))
+        _produce_test(locals(), gen_test_min_type, (op,))
+
+    for i, op in enumerate(UNARY_OPERATORS):
+        _produce_test(locals(), gen_test_uniop, (op,))
+        _produce_test(locals(), gen_test_min_type_unary, (op,))
+
+
+class ExpressionTest(unittest.TestCase):
+    def gen_test_expression(expression_str, expected_type, use_py2_division=False):
+        def test_expression(self):
+            expr = tree(expression_str).children[0].children[0]
+            inferred_type = numeric_expr_type(
+                expr, use_py2_division, type_for_unknown=InferredType.INT_OR_FLOAT
+            )
+            self.assertEqual(expected_type, inferred_type)
+
+        return test_expression
+
+    for expr, expected in SAMPLE_EXPRESSIONS:
+        _produce_test(locals(), gen_test_expression, (expr + "\n", expected))
+    for expr, expected in SAMPLE_PY2_EXPRESSIONS:
+        _produce_test(locals(), gen_test_expression, (expr + "\n", expected, True))