docs/source/tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "====================\n",
    "Parsing and Visiting\n",
    "====================\n",
    "\n",
    "LibCST provides helpers to parse source code string as concrete syntax tree. In order to perform static analysis to identify patterns in the tree or modify the tree programmatically, we can use visitor pattern to traverse the tree. In this tutorial, we demonstrate a common three-step-workflow to build an automated refactoring (codemod) application:\n",
    "\n",
    "1. `Parse Source Code <#Parse-Source-Code>`_\n",
    "2. `Build Visitor or Transformer <#Build-Visitor-or-Transformer>`_\n",
    "3. `Generate Source Code <#Generate-Source-Code>`_\n",
    "\n",
    "Parse Source Code\n",
    "=================\n",
    "LibCST provides various helpers to parse source code as concrete syntax tree: :func:`~libcst.parse_module`, :func:`~libcst.parse_expression` and :func:`~libcst.parse_statement` (see :doc:`Parsing <parser>` for more detail). The default :class:`~libcst.CSTNode` repr provides pretty print formatting for reading the tree easily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"../../\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import libcst as cst\n",
    "\n",
    "cst.parse_expression(\"1 + 2\")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Example: add typing annotation from pyi stub file to Python source\n",
    "------------------------------------------------------------------\n",
    "Python `typing annotation <https://mypy.readthedocs.io/en/latest/cheat_sheet_py3.html>`_ was added in Python 3.5. Some Python applications add typing annotations in separate ``pyi`` stub files in order to support old Python versions. When applications decide to stop supporting old Python versions, they'll want to automatically copy the type annotation from a pyi file to a source file. Here we demonstrate how to do that easliy using LibCST. The first step is to parse the pyi stub and source files as trees."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "py_source = '''\n",
    "class PythonToken(Token):\n",
    "    def __repr__(self):\n",
    "        return ('TokenInfo(type=%s, string=%r, start_pos=%r, prefix=%r)' %\n",
    "                self._replace(type=self.type.name))\n",
    "\n",
    "def tokenize(code, version_info, start_pos=(1, 0)):\n",
    "    \"\"\"Generate tokens from a the source code (string).\"\"\"\n",
    "    lines = split_lines(code, keepends=True)\n",
    "    return tokenize_lines(lines, version_info, start_pos=start_pos)\n",
    "'''\n",
    "\n",
    "pyi_source = '''\n",
    "class PythonToken(Token):\n",
    "    def __repr__(self) -> str: ...\n",
    "\n",
    "def tokenize(\n",
    "    code: str, version_info: PythonVersionInfo, start_pos: Tuple[int, int] = (1, 0)\n",
    ") -> Generator[PythonToken, None, None]: ...\n",
    "'''\n",
    "\n",
    "source_tree = cst.parse_module(py_source)\n",
    "stub_tree = cst.parse_module(pyi_source)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Build Visitor or Transformer\n",
    "============================\n",
    "For traversing and modifying the tree, LibCST provides Visitor and Transformer classes similar to the `ast module <https://docs.python.org/3/library/ast.html#ast.NodeVisitor>`_. To implement a visitor (read only) or transformer (read/write), simply implement a subclass of :class:`~libcst.CSTVisitor` or :class:`~libcst.CSTTransformer` (see :doc:`Visitors <visitors>` for more detail).\n",
    "In the typing example, we need to implement a visitor to collect typing annotation from the stub tree and a transformer to copy the annotation to the function signature. In the visitor, we implement ``visit_FunctionDef`` to collect annotations. Later in the transformer, we implement ``leave_FunctionDef`` to add the collected annotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Tuple, Dict, Optional\n",
    "\n",
    "\n",
    "class TypingCollector(cst.CSTVisitor):\n",
    "    def __init__(self):\n",
    "        # stack for storing the canonical name of the current function\n",
    "        self.stack: List[Tuple[str, ...]] = []\n",
    "        # store the annotations\n",
    "        self.annotations: Dict[\n",
    "            Tuple[str, ...],  # key: tuple of canonical class/function name\n",
    "            Tuple[cst.Parameters, Optional[cst.Annotation]],  # value: (params, returns)\n",
    "        ] = {}\n",
    "\n",
    "    def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:\n",
    "        self.stack.append(node.name.value)\n",
    "\n",
    "    def leave_ClassDef(self, node: cst.ClassDef) -> None:\n",
    "        self.stack.pop()\n",
    "\n",
    "    def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:\n",
    "        self.stack.append(node.name.value)\n",
    "        self.annotations[tuple(self.stack)] = (node.params, node.returns)\n",
    "        return (\n",
    "            False\n",
    "        )  # pyi files don't support inner functions, return False to stop the traversal.\n",
    "\n",
    "    def leave_FunctionDef(self, node: cst.FunctionDef) -> None:\n",
    "        self.stack.pop()\n",
    "\n",
    "\n",
    "class TypingTransformer(cst.CSTTransformer):\n",
    "    def __init__(self, annotations):\n",
    "        # stack for storing the canonical name of the current function\n",
    "        self.stack: List[Tuple[str, ...]] = []\n",
    "        # store the annotations\n",
    "        self.annotations: Dict[\n",
    "            Tuple[str, ...],  # key: tuple of canonical class/function name\n",
    "            Tuple[cst.Parameters, Optional[cst.Annotation]],  # value: (params, returns)\n",
    "        ] = annotations\n",
    "\n",
    "    def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:\n",
    "        self.stack.append(node.name.value)\n",
    "\n",
    "    def leave_ClassDef(\n",
    "        self, original_node: cst.ClassDef, updated_node: cst.ClassDef\n",
    "    ) -> cst.CSTNode:\n",
    "        self.stack.pop()\n",
    "        return updated_node\n",
    "\n",
    "    def visit_FunctionDef(self, node: cst.FunctionDef) -> Optional[bool]:\n",
    "        self.stack.append(node.name.value)\n",
    "        return (\n",
    "            False\n",
    "        )  # pyi files don't support inner functions, return False to stop the traversal.\n",
    "\n",
    "    def leave_FunctionDef(\n",
    "        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef\n",
    "    ) -> cst.CSTNode:\n",
    "        key = tuple(self.stack)\n",
    "        self.stack.pop()\n",
    "        if key in self.annotations:\n",
    "            annotations = self.annotations[key]\n",
    "            return updated_node.with_changes(\n",
    "                params=annotations[0], returns=annotations[1]\n",
    "            )\n",
    "        return updated_node\n",
    "\n",
    "\n",
    "visitor = TypingCollector()\n",
    "stub_tree.visit(visitor)\n",
    "transformer = TypingTransformer(visitor.annotations)\n",
    "modified_tree = source_tree.visit(transformer)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "Generate Source Code\n",
    "====================\n",
    "Generating the source code from a cst tree is as easy as accessing the :attr:`~libcst.Module.code` attribute on :class:`~libcst.Module`. After the code generation, we often use `ufmt <https://ufmt.omnilib.dev/en/stable/>`_ to reformate the code to keep a consistent coding style."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(modified_tree.code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use difflib to show the changes to verify type annotations were added as expected.\n",
    "import difflib\n",
    "\n",
    "print(\n",
    "    \"\".join(\n",
    "        difflib.unified_diff(py_source.splitlines(1), modified_tree.code.splitlines(1))\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "raw_mimetype": "text/restructuredtext"
   },
   "source": [
    "For the sake of efficiency, we don't want to re-write the file when the transformer doesn't change the source code. We can use :meth:`~libcst.CSTNode.deep_equals` to check whether two trees have the same source code. Note that ``==`` checks the identity of tree object instead of representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not modified_tree.deep_equals(source_tree):\n",
    "    ...  # write to file"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}