[Paper Review] TypeT5: Seq2seq Type Inference using Static Analysis

Jiayi Wei, Greg Durrett+1 more|arXiv (Cornell University)|Mar 16, 2023

Software Engineering Research4 citations

TL;DR

TypeT5 proposes a seq2seq type inference method for Python that leverages CodeT5 with static analysis to construct global code contexts and an iterative decoding scheme for improved type prediction. It achieves state-of-the-art accuracy, especially on rare and complex types, and reduces user correction effort to one in five predictions in interactive settings.

ABSTRACT

There has been growing interest in automatically predicting missing type annotations in programs written in Python and JavaScript. While prior methods have achieved impressive accuracy when predicting the most common types, they often perform poorly on rare or complex types. In this paper, we present a new type inference method that treats type prediction as a code infilling task by leveraging CodeT5, a state-of-the-art seq2seq pre-trained language model for code. Our method uses static analysis to construct dynamic contexts for each code element whose type signature is to be predicted by the model. We also propose an iterative decoding scheme that incorporates previous type predictions in the model's input context, allowing information exchange between related code elements. Our evaluation shows that the proposed approach, TypeT5, not only achieves a higher overall accuracy (particularly on rare and complex types) but also produces more coherent results with fewer type errors -- while enabling easy user intervention.

Motivation & Objective

• To improve type inference in untyped Python code by treating type prediction as a code infilling task using a pre-trained seq2seq model.
• To enhance model performance on rare and complex types by incorporating non-local code context through static analysis.
• To enable coherent type assignments by propagating predictions across related code elements using iterative decoding.
• To support interactive development by allowing users to efficiently correct model predictions, reducing annotation overhead.
• To demonstrate that combining pre-trained models with static analysis yields superior results over prior learning-based and baseline approaches.

Proposed method

• TypeT5 treats type inference as a code infilling task, using CodeT5 to generate type annotations conditioned on source code and dynamically constructed context.
• It constructs a usage graph via static analysis to identify relevant code elements (users and usees) for each target code element, forming extended contextual inputs.
• The model encodes these contextually enriched inputs using subword tokenization and BPE, enabling generation of arbitrary type expressions including parametric and user-defined types.
• An iterative decoding scheme is introduced, where previously predicted types are injected back into the input context to enable information exchange between related code elements.
• The system fine-tunes CodeT5 on untyped Python code with ground-truth type annotations, using a masked language modeling objective for sequence-to-sequence learning.
• The approach supports interactive use by allowing users to inspect and correct predictions, with the model re-informing context in subsequent predictions.

Experimental results

Research questions

• RQ1 Can a seq2seq pre-trained model like CodeT5 be effectively adapted for type inference in untyped Python code when augmented with static analysis?
• RQ2 Does incorporating non-local code context via usage graphs improve prediction accuracy, especially for rare and complex types?
• RQ3 Can iterative decoding that propagates type predictions across code elements enhance the coherence and correctness of generated types?
• RQ4 How does TypeT5 compare to state-of-the-art type inference tools in terms of accuracy and robustness on challenging type patterns?
• RQ5 To what extent can interactive correction reduce the human effort required to fully annotate an untyped codebase?

Key findings

• TypeT5 outperforms three state-of-the-art type inference tools and a CodeT5 baseline, achieving significantly higher overall accuracy, especially on rare and complex types.
• The use of static analysis to build context from usage graphs improves performance on challenging types, such as nested parametric types like dict[int, list[PythonType]] and user-defined types.
• The iterative decoding scheme reduces type errors by improving coherence, as evidenced by fewer type constraint violations reported by the type checker.
• In an interactive setting, users need to correct only one in every five model predictions to fully annotate an untyped codebase, indicating high practical usability.
• Ablation studies confirm that both static analysis context and iterative decoding are critical components for the model’s performance gains.
• TypeT5 is the first learning-based method capable of predicting both parametric and user-defined types using a seq2seq framework.