  Jupyter notebooks enable developers to interleave code snippets with
rich-text and in-line visualizations. Data scientists use Jupyter notebook as
the de-facto standard for creating and sharing machine-learning based
solutions, primarily written in Python. Recent studies have demonstrated,
however, that a large portion of Jupyter notebooks available on public
platforms are undocumented and lacks a narrative structure. This reduces the
readability of these notebooks. To address this shortcoming, this paper
presents HeaderGen, a novel tool-based approach that automatically annotates
code cells with categorical markdown headers based on a taxonomy of ML
operations, and classifies and displays function calls according to this
taxonomy. For this functionality to be realized, HeaderGen enhances an existing
call graph analysis in PyCG. To improve precision, HeaderGen extends PyCG's
analysis with support for handling external library code and flow-sensitivity.
The former is realized by facilitating the resolution of function return-types.
The evaluation on 15 real-world Jupyter notebooks from Kaggle shows that
HeaderGen's underlying call graph analysis yields high accuracy (95.6%
precision and 95.3% recall). This is because HeaderGen can resolve return-types
of external libraries where existing type inference tools such as pytype (by
Google), pyright (by Microsoft), and Jedi fall short. The header generation has
a precision of 85.7% and a recall rate of 92.8%. In a user study, HeaderGen
helps participants finish comprehension and navigation tasks faster. To further
evaluate the type inference capability of tools, we introduce TypeEvalPy, a
framework for evaluating type inference tools with a micro-benchmark containing
154 code snippets and 845 type annotations. Our comparative analysis on four
tools revealed that HeaderGen outperforms other tools in exact matches with the
ground truth.
