AST - Get assignment types.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "code = \"\"\"\\\n",
    "a = 1\n",
    "b = 2.0\n",
    "c = \"hello world\"\n",
    "\"\"\"\n",
    "\n",
    "code_obj = compile(code, \"<string>\", \"exec\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "module 'ast' has no attribute '_fields'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-2-c11a8f80af1a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0ma\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparse\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mnode\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwalk\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mast\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      5\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.8/ast.py\u001b[0m in \u001b[0;36mwalk\u001b[0;34m(node)\u001b[0m\n\u001b[1;32m    330\u001b[0m     \u001b[0;32mwhile\u001b[0m \u001b[0mtodo\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    331\u001b[0m         \u001b[0mnode\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtodo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpopleft\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 332\u001b[0;31m         \u001b[0mtodo\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mextend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miter_child_nodes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    333\u001b[0m         \u001b[0;32myield\u001b[0m \u001b[0mnode\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    334\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.8/ast.py\u001b[0m in \u001b[0;36miter_child_nodes\u001b[0;34m(node)\u001b[0m\n\u001b[1;32m    214\u001b[0m     \u001b[0;32mand\u001b[0m \u001b[0mall\u001b[0m \u001b[0mitems\u001b[0m \u001b[0mof\u001b[0m \u001b[0mfields\u001b[0m \u001b[0mthat\u001b[0m \u001b[0mare\u001b[0m \u001b[0mlists\u001b[0m \u001b[0mof\u001b[0m \u001b[0mnodes\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    215\u001b[0m     \"\"\"\n\u001b[0;32m--> 216\u001b[0;31m     \u001b[0;32mfor\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfield\u001b[0m \u001b[0;32min\u001b[0m \u001b[0miter_fields\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    217\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfield\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mAST\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    218\u001b[0m             \u001b[0;32myield\u001b[0m \u001b[0mfield\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/lib/python3.8/ast.py\u001b[0m in \u001b[0;36miter_fields\u001b[0;34m(node)\u001b[0m\n\u001b[1;32m    202\u001b[0m     \u001b[0mthat\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0mpresent\u001b[0m \u001b[0mon\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    203\u001b[0m     \"\"\"\n\u001b[0;32m--> 204\u001b[0;31m     \u001b[0;32mfor\u001b[0m \u001b[0mfield\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mnode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_fields\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    205\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    206\u001b[0m             \u001b[0;32myield\u001b[0m \u001b[0mfield\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfield\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAttributeError\u001b[0m: module 'ast' has no attribute '_fields'"
     ]
    }
   ],
   "source": [
    "import ast\n",
    "\n",
    "a = ast.parse(code)\n",
    "for node in ast.walk(ast):\n",
    "    print(node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ast import *\n",
    "\n",
    "sorted(dir())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<_ast.Module object at 0x7ff6e432f040>\n",
      "<_ast.Assign object at 0x7ff6e432f070>\n",
      "<_ast.Assign object at 0x7ff6e432f1c0>\n",
      "<_ast.Assign object at 0x7ff6e432f250>\n",
      "<_ast.Assign object at 0x7ff6e432f2e0>\n",
      "<_ast.Assign object at 0x7ff6e432f3a0>\n",
      "<_ast.Expr object at 0x7ff6e432f4c0>\n",
      "<_ast.Expr object at 0x7ff6e432f5b0>\n",
      "<_ast.Name object at 0x7ff6e432f0a0>\n",
      "<_ast.Call object at 0x7ff6e432f0d0>\n",
      "<_ast.Name object at 0x7ff6e432f1f0>\n",
      "<_ast.Constant object at 0x7ff6e432f220>\n",
      "<_ast.Name object at 0x7ff6e432f280>\n",
      "<_ast.Constant object at 0x7ff6e432f2b0>\n",
      "<_ast.Name object at 0x7ff6e432f310>\n",
      "<_ast.Call object at 0x7ff6e432f340>\n",
      "<_ast.Name object at 0x7ff6e432f3d0>\n",
      "<_ast.Call object at 0x7ff6e432f430>\n",
      "<_ast.Call object at 0x7ff6e432f4f0>\n",
      "<_ast.Call object at 0x7ff6e432f5e0>\n",
      "<_ast.Store object at 0x7ff6e8b94a00>\n",
      "<_ast.Name object at 0x7ff6e432f100>\n",
      "<_ast.Constant object at 0x7ff6e432f160>\n",
      "<_ast.Constant object at 0x7ff6e432f190>\n",
      "<_ast.Store object at 0x7ff6e8b94a00>\n",
      "<_ast.Store object at 0x7ff6e8b94a00>\n",
      "<_ast.Store object at 0x7ff6e8b94a00>\n",
      "<_ast.Name object at 0x7ff6e432f370>\n",
      "<_ast.Store object at 0x7ff6e8b94a00>\n",
      "<_ast.Name object at 0x7ff6e432f460>\n",
      "<_ast.Constant object at 0x7ff6e432f490>\n",
      "<_ast.Name object at 0x7ff6e432f520>\n",
      "<_ast.Name object at 0x7ff6e432f550>\n",
      "<_ast.Name object at 0x7ff6e432f610>\n",
      "<_ast.Name object at 0x7ff6e432f640>\n",
      "<_ast.Load object at 0x7ff6e8b949a0>\n",
      "<_ast.Load object at 0x7ff6e8b949a0>\n",
      "<_ast.Load object at 0x7ff6e8b949a0>\n",
      "<_ast.Load object at 0x7ff6e8b949a0>\n",
      "<_ast.Load object at 0x7ff6e8b949a0>\n",
      "<_ast.Load object at 0x7ff6e8b949a0>\n",
      "<_ast.Load object at 0x7ff6e8b949a0>\n",
      "Found variable 'myRangeVar' with a value of 'range(1, 10)' (type: 'Call')\n",
      "Found variable 'myIntVar' with a value of '1' (type: 'int')\n",
      "Found variable 'myStrVar' with a value of 'hello world' (type: 'str')\n",
      "Found variable 'myTest' with a value of 'fakeFunct()' (type: 'Call')\n",
      "Found variable 'myTestTwo' with a value of 'fakeFunct(20)' (type: 'Call')\n"
     ]
    }
   ],
   "source": [
    "import ast\n",
    "\n",
    "class ASTExplorer:\n",
    "    def __init__(self, source):\n",
    "        self.tree = ast.parse(source, mode=\"exec\")\n",
    "        self.result = list()\n",
    "\n",
    "    class ASTResult:\n",
    "        def __init__(self, var, expr, vType):\n",
    "            self.var = var\n",
    "            self.expression = expr\n",
    "            self.vType = vType\n",
    "\n",
    "    def _getLineAssignment(self, lineno):\n",
    "        return next((node for node in ast.walk(self.tree) if isinstance(node, ast.Name) and node.lineno == lineno), None)\n",
    "\n",
    "    def getVariables(self):\n",
    "        for node in ast.walk(self.tree):\n",
    "            print(node)\n",
    "            if not isinstance(node, ast.Assign):\n",
    "                continue\n",
    "            nodeValue = node.value\n",
    "            nodeVariable = self._getLineAssignment(node.lineno).id\n",
    "            if(isinstance(nodeValue, ast.Constant)):\n",
    "                nodeExpression = node.value.value\n",
    "                self.result.append(self.ASTResult(nodeVariable, nodeExpression, type(nodeExpression)))\n",
    "                continue\n",
    "            elif(isinstance(nodeValue, ast.Call)):\n",
    "                callFunc = nodeValue.func.id\n",
    "                callArgs = \"(\" + (\", \".join([str(x.value) for x in nodeValue.args])) + \")\"\n",
    "                self.result.append(self.ASTResult(nodeVariable, f\"{callFunc}{callArgs}\", ast.Call))\n",
    "            #elif... other type handling\n",
    "        return self.result\n",
    "    \n",
    "\n",
    "source_code = \"\"\"\\\n",
    "myRangeVar = range(1, 10)\n",
    "myIntVar = 1\n",
    "myStrVar = \"hello world\"\n",
    "myTest = fakeFunct()\n",
    "myTestTwo = fakeFunct(20)\n",
    "print(a)\n",
    "print(b)\n",
    "\"\"\"\n",
    "\n",
    "explorer = ASTExplorer(source_code)\n",
    "for result in explorer.getVariables():\n",
    "    print(f\"Found variable '{result.var}' with a value of '{result.expression}' (type: '{result.vType.__name__}')\")"
   ]
  }
 ],
 "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.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}