nbs/09_magics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# magics\n",
    "\n",
    "> IPython magics for nbstata\n",
    "- order: 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp magics\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "from nbstata.config import Config\n",
    "from nbstata.misc_utils import print_red\n",
    "from nbstata.stata import obs_count, macro_expand, get_global\n",
    "from nbstata.stata_session import warn_re_unclosed_comment_block_if_needed\n",
    "import nbstata.browse as browse\n",
    "from fastcore.basics import patch_to\n",
    "import re\n",
    "import urllib\n",
    "from pkg_resources import resource_filename\n",
    "from bs4 import BeautifulSoup as bs\n",
    "import configparser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def print_kernel(msg, kernel):\n",
    "    msg = re.sub(r'$', r'\\r\\n', msg, flags=re.MULTILINE)\n",
    "    msg = re.sub(r'[\\r\\n]{1,2}[\\r\\n]{1,2}', r'\\r\\n', msg, flags=re.MULTILINE)\n",
    "    stream_content = {'text': msg, 'name': 'stdout'}\n",
    "    kernel.send_response(kernel.iopub_socket, 'stream', stream_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _construct_abbrev_dict():\n",
    "    def _all_abbrevs(full_command, shortest_abbrev):\n",
    "        for j in range(len(shortest_abbrev), len(full_command)):\n",
    "            yield full_command[0:j]\n",
    "    out = {}\n",
    "    abbrev_list = [\n",
    "        ('browse', 'br'),\n",
    "        ('frbrowse', 'frbr'),\n",
    "        ('help', 'h'),\n",
    "        ('quietly', 'q'),\n",
    "    ]\n",
    "    for full_command, shortest_abbrev in abbrev_list:\n",
    "        out.update(\n",
    "            {abbrev: full_command\n",
    "             for abbrev in _all_abbrevs(full_command, shortest_abbrev)}\n",
    "        )\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'br': 'browse',\n",
       " 'bro': 'browse',\n",
       " 'brow': 'browse',\n",
       " 'brows': 'browse',\n",
       " 'frbr': 'frbrowse',\n",
       " 'frbro': 'frbrowse',\n",
       " 'frbrow': 'frbrowse',\n",
       " 'frbrows': 'frbrowse',\n",
       " 'h': 'help',\n",
       " 'he': 'help',\n",
       " 'hel': 'help',\n",
       " 'q': 'quietly',\n",
       " 'qu': 'quietly',\n",
       " 'qui': 'quietly',\n",
       " 'quie': 'quietly',\n",
       " 'quiet': 'quietly',\n",
       " 'quietl': 'quietly'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#| hide\n",
    "_construct_abbrev_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class StataMagics():\n",
    "    \"\"\"Class for handling magics\"\"\"\n",
    "    magic_regex = re.compile(\n",
    "        r'\\A\\*?(?P<magic>%%?\\w+?)(?P<code>[\\s,]+.*?)?\\Z', flags=re.DOTALL + re.MULTILINE)\n",
    "\n",
    "    # Format: magic_name: help_content\n",
    "    available_magics = {\n",
    "        '%browse': '{} [-h] [varlist] [if] [in] [, nolabel noformat]',\n",
    "        '%head': '{} [-h] [N] [varlist] [if] [, nolabel noformat]',\n",
    "        '%tail': '{} [-h] [N] [varlist] [if] [, nolabel noformat]',\n",
    "        '%frbrowse': '{} [-h] framename: [varlist] [if] [in] [, nolabel noformat]',\n",
    "        '%frhead': '{} [-h] framename: [N] [varlist] [if] [, nolabel noformat]',\n",
    "        '%frtail': '{} [-h] framename: [N] [varlist] [if] [, nolabel noformat]',\n",
    "        '%locals': '',\n",
    "        '%delimit': '',\n",
    "        '%help': '{} [-h] command_or_topic_name',\n",
    "        '%set': '{} [-h] key = value',\n",
    "        '%%set': '{} [-h]\\nkey1 = value1\\n[key2 = value2]\\n[...]',\n",
    "        '%status': '',\n",
    "        '%%quietly': '',\n",
    "        '%%noecho': '',\n",
    "        '%%echo': '',\n",
    "    }\n",
    "    \n",
    "    abbrev_dict = _construct_abbrev_dict()\n",
    "    \n",
    "    csshelp_default = resource_filename(\n",
    "        'nbstata', 'css/_StataKernelHelpDefault.css'\n",
    "    )\n",
    "\n",
    "    def magic_quietly(self, code, kernel, cell):\n",
    "        \"\"\"Suppress all display for the current cell.\"\"\"\n",
    "        cell.quietly = True\n",
    "        return code\n",
    "\n",
    "    def magic_noecho(self, code, kernel, cell):\n",
    "        \"\"\"Suppress echo for the current cell.\"\"\"\n",
    "        cell.noecho = True\n",
    "        cell.echo = False\n",
    "        return code\n",
    "    \n",
    "    def magic_echo(self, code, kernel, cell):\n",
    "        \"\"\"Suppress echo for the current cell.\"\"\"\n",
    "        cell.noecho = False\n",
    "        cell.echo = True\n",
    "        return code\n",
    "    \n",
    "    def magic_delimit(self, code, kernel, cell):\n",
    "        delim = ';' if kernel.stata_session.sc_delimiter else 'cr'\n",
    "        print_kernel(f'Current Stata command delimiter: {delim}', kernel)\n",
    "        return ''\n",
    "    \n",
    "    def magic_status(self, code, kernel, cell):\n",
    "        kernel.nbstata_config.display_status()\n",
    "        return ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def _unabbrev_magic_name(self, raw_name):\n",
    "    last_percent = raw_name.rfind('%')\n",
    "    percent_part = raw_name[:last_percent+1]\n",
    "    raw_name_part = raw_name[last_percent+1:]\n",
    "    if raw_name_part in self.abbrev_dict:\n",
    "        name_part = self.abbrev_dict[raw_name_part]\n",
    "    else:\n",
    "        name_part = raw_name_part\n",
    "    return percent_part + name_part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from fastcore.test import test_eq, test_fail"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_instance = StataMagics()\n",
    "test_eq(test_instance._unabbrev_magic_name(\"*%non_name\"), \"*%non_name\")\n",
    "test_eq(test_instance._unabbrev_magic_name(\"%%se\"), \"%%se\")\n",
    "test_eq(test_instance._unabbrev_magic_name(\"%br\"), \"%browse\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _parse_magic_name_code(match):\n",
    "    v = match.groupdict()\n",
    "    for k in v:\n",
    "        v[k] = v[k] if v[k] is not None else ''                \n",
    "    name = v['magic'].strip()\n",
    "    code = v['code'].strip()\n",
    "    return name, code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def _parse_code_for_magic(self, code):\n",
    "    match = self.magic_regex.match(code.strip())\n",
    "    if match:\n",
    "        raw_name, mcode = _parse_magic_name_code(match)\n",
    "        name = self._unabbrev_magic_name(raw_name)\n",
    "        if name in {'%quietly', '%noecho', '%echo'}:\n",
    "            print_red(\n",
    "                f\"Warning: The correct syntax for a cell magic is '%{name}', not '{name}'. \"\n",
    "                \"In v1.0, nbstata may trigger an error instead of just a warning.\"\n",
    "            )\n",
    "            name = '%' + name\n",
    "        elif name == \"%set\" and len(code.splitlines()) > 1:\n",
    "            print_red(\n",
    "                f\"Warning: The correct syntax for the multi-line 'set' magic is '%{name}', not '{name}'. \"\n",
    "                \"In v1.0, nbstata may trigger an error instead of just a warning.\"\n",
    "            )\n",
    "            name = '%' + name\n",
    "        elif name == \"%%set\" and len(code.splitlines()) == 1:\n",
    "            print_red(\n",
    "                f\"Warning: The correct syntax for the single-line 'set' magic is '%set', not '{name}'. \"\n",
    "                \"In v1.0, nbstata may trigger an error instead of just a warning.\"\n",
    "            )\n",
    "            name = '%set'\n",
    "        elif name not in self.available_magics:\n",
    "            raise ValueError(f\"Unknown magic {name}.\")\n",
    "        return name, mcode\n",
    "    else:\n",
    "        return None, code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "---\n",
       "\n",
       "[source](https://github.com/hugetim/nbstata/blob/master/nbstata/magics.py#L125){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### StataMagics._parse_code_for_magic\n",
       "\n",
       ">      StataMagics._parse_code_for_magic (code)"
      ],
      "text/plain": [
       "---\n",
       "\n",
       "[source](https://github.com/hugetim/nbstata/blob/master/nbstata/magics.py#L125){target=\"_blank\" style=\"float:right; font-size:smaller\"}\n",
       "\n",
       "### StataMagics._parse_code_for_magic\n",
       "\n",
       ">      StataMagics._parse_code_for_magic (code)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "show_doc(StataMagics._parse_code_for_magic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_instance = StataMagics()\n",
    "test_eq(test_instance._parse_code_for_magic(\"%browse\"), ('%browse', \"\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%browse\"), ('%browse', \"\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%br\"), ('%browse', \"\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%browse,\"), ('%browse', \",\"))\n",
    "test_fail(test_instance._parse_code_for_magic, args=(\"*%blah\\nreg var1\"))\n",
    "test_fail(test_instance._parse_code_for_magic, args=(\"*%se echo=True\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mWarning: The correct syntax for a cell magic is '%%echo', not '%echo'. In v1.0, nbstata may trigger an error instead of just a warning.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%browse, noformat\"), ('%browse', \", noformat\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%head 1, noformat\"), ('%head', \"1, noformat\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"%browse -h\"), ('%browse', \"-h\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%browse var1, nolabel\"), ('%browse', \"var1, nolabel\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%help reg\"), ('%help', \"reg\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%h reg\"), ('%help', \"reg\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%%echo\\nreg var1\"), ('%%echo', \"reg var1\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%echo\\nreg var1\"), ('%%echo', \"reg var1\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%set echo = True\"), ('%set', \"echo = True\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%%set\\necho = True\"), ('%%set', \"echo = True\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mWarning: The correct syntax for the single-line 'set' magic is '%set', not '%%set'. In v1.0, nbstata may trigger an error instead of just a warning.\u001b[0m\n",
      "\u001b[31mWarning: The correct syntax for the multi-line 'set' magic is '%%set', not '%set'. In v1.0, nbstata may trigger an error instead of just a warning.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%%set echo = True\"), ('%set', \"echo = True\"))\n",
    "test_eq(test_instance._parse_code_for_magic(\"*%set\\necho = True\"), ('%%set', \"echo = True\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('%browse', 'in 1/10')"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_instance._parse_code_for_magic(\"*%browse in 1/10\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def _do_magic(self, name, code, kernel, cell):\n",
    "    if code.startswith('-h') or code.startswith('--help') or (name == \"%help\" and (not code or code.isspace())):\n",
    "        print_kernel(self.available_magics[name].format(name), kernel)\n",
    "        return ''\n",
    "    else:\n",
    "        return getattr(self, \"magic_\" + name.lstrip('%'))(code, kernel, cell)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "''.isspace()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic(self, code, kernel, cell):\n",
    "    try:\n",
    "        if not kernel.ipydatagrid_height_set:\n",
    "            browse.set_ipydatagrid_height()\n",
    "            kernel.ipydatagrid_height_set = True\n",
    "        name, code = self._parse_code_for_magic(code)\n",
    "    except ValueError as e:\n",
    "        print_kernel(str(e), kernel)\n",
    "    else:\n",
    "        if name:\n",
    "            code = self._do_magic(name, code, kernel, cell)\n",
    "    return code        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _formatted_local_list(local_dict):\n",
    "    std_len = 14\n",
    "    str_reps = []\n",
    "    for n in local_dict:\n",
    "        if len(n) <= std_len:\n",
    "            str_reps.append(f\"{n}:{' '*(std_len-len(n))} {local_dict[n]}\")\n",
    "        else:\n",
    "            str_reps.append(f\"{n}:\\n{' '*std_len}  {local_dict[n]}\")\n",
    "    return \"\\n\".join(str_reps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test1:          blah blah\n",
      "long_name_local:\n",
      "                s\n",
      "local2:         1234\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "fake_locals = {'test1': \"blah blah\", 'long_name_local': 's', 'local2': 1234}\n",
    "print(_formatted_local_list(fake_locals))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "print(_formatted_local_list({}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_locals(self, code, kernel, cell):\n",
    "    local_dict = kernel.stata_session.get_local_dict()\n",
    "    print_kernel(_formatted_local_list(local_dict), kernel)\n",
    "    return ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#| hide\n",
    "* [https://docs.python.org/3.10/library/configparser.html#customizing-parser-behaviour](https://docs.python.org/3.10/library/configparser.html#customizing-parser-behaviour)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _get_new_settings(code):\n",
    "    parser = configparser.ConfigParser(\n",
    "        empty_lines_in_values=False,\n",
    "        comment_prefixes=('*','//', '/*'), # '/*': to not cause error when commenting out for Stata purposes only\n",
    "        inline_comment_prefixes=('//',),\n",
    "    )\n",
    "    parser.read_string(\"[set]\\n\" + code.strip(), source=\"set\")\n",
    "    return dict(parser.items('set'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'echo': 'True'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_get_new_settings(\"echo = True\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'missing': 'NOTHING', 'echo': 'False'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_get_new_settings(\"\"\"\n",
    "/*\n",
    "missing = NOTHING\n",
    "echo = False\n",
    "*/\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#| hide\n",
    "_get_new_settings(\"/* */ echo = True\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _clean_error_message(err_str):\n",
    "    return (err_str\n",
    "            .replace(\"While reading from 'set'\", \"\")\n",
    "            .replace(\" in section 'set' already exists\", \" already set above\")\n",
    "            .replace(\"Source contains \", \"\")\n",
    "            .replace(\" 'set'\\n\", \"\\n\")\n",
    "           )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _process_new_settings(settings, kernel):\n",
    "    kernel.nbstata_config.update(settings)\n",
    "    kernel.nbstata_config.update_graph_config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_set(self, code, kernel, cell):\n",
    "    try:\n",
    "        settings = _get_new_settings(code)\n",
    "    except configparser.Error as err:\n",
    "        print_red(f\"set error:\\n    {_clean_error_message(str(err))}\")\n",
    "    else:\n",
    "        _process_new_settings(settings, kernel)\n",
    "        warn_re_unclosed_comment_block_if_needed(code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mset error:\n",
      "     [line  4]: option 'echo' already set above\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "code = \"\"\"\n",
    "/*\n",
    "echo = None\n",
    "echo = False\n",
    "*/\"\"\"\n",
    "try:\n",
    "    settings = _get_new_settings(code)\n",
    "except configparser.Error as err:\n",
    "    print_red(f\"set error:\\n    {_clean_error_message(str(err))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mset error:\n",
      "    parsing errors:\n",
      "\t[line  2]: 'echo False'\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "code = \"\"\"echo False\"\"\"\n",
    "try:\n",
    "    settings = _get_new_settings(code)\n",
    "except configparser.Error as err:\n",
    "    print_red(f\"set error:\\n    {_clean_error_message(str(err))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mset error:\n",
      "    parsing errors:\n",
      "\t[line  2]: 'echo False\\n'\n",
      "\t[line  3]: 'missing No'\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "#| hide\n",
    "code = \"\"\"echo False\n",
    "missing No\"\"\"\n",
    "try:\n",
    "    settings = _get_new_settings(code)\n",
    "except configparser.Error as err:\n",
    "    print_red(f\"set error:\\n    {_clean_error_message(str(err))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Browse magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_browse(self, code, kernel, cell):\n",
    "    \"\"\"Display data interactively.\"\"\"\n",
    "    try:\n",
    "        expanded_code = macro_expand(code)\n",
    "        params = browse.browse_df_params(\n",
    "            expanded_code, obs_count(), kernel.nbstata_config.env['missing'],\n",
    "        )\n",
    "        sformat = params[-1]\n",
    "        df = browse.get_df(*params)\n",
    "        browse.display_df_as_ipydatagrid(df)\n",
    "    except Exception as e:\n",
    "        print_kernel(f\"browse failed.\\r\\n{e}\", kernel)\n",
    "    return ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class Frame():\n",
    "    \"\"\"Class for generating Stata select_var for getAsDict\"\"\"\n",
    "    def __init__(self, framename):\n",
    "        self.original_framename = get_global('c(frame)')\n",
    "        self.framename = framename\n",
    "            \n",
    "    def __enter__(self):\n",
    "        import sfi\n",
    "        try:\n",
    "            frame = sfi.Frame.connect(self.framename)\n",
    "        except sfi.FrameError:\n",
    "            raise ValueError(f\"frame {self.framename} not found\")\n",
    "        else:\n",
    "            frame.changeToCWF()\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, exc_tb):\n",
    "        import sfi\n",
    "        orig_frame = sfi.Frame.connect(self.original_framename)\n",
    "        orig_frame.changeToCWF()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _parse_frame_prefix(code):\n",
    "    pattern = re.compile(\n",
    "        r'\\A(?P<frame>\\w+)[ \\t]*(?:\\:[ \\t]*(?P<code>.*?))?\\Z', flags=re.DOTALL)\n",
    "    match = pattern.match(code)\n",
    "    if not match:\n",
    "        raise ValueError(\"invalid syntax: missing framename or colon?\")\n",
    "    v = match.groupdict()\n",
    "    for k in v:\n",
    "        v[k] = v[k] if v[k] is not None else ''                \n",
    "    framename = v['frame'].strip()\n",
    "    main_code = v['code'].strip()\n",
    "    return framename, main_code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "test_eq(_parse_frame_prefix('subtask: if var1==1'), ('subtask', 'if var1==1'))\n",
    "test_eq(_parse_frame_prefix('subtask'), ('subtask', ''))\n",
    "test_eq(_parse_frame_prefix('subtask:'), ('subtask', ''))\n",
    "test_fail(_parse_frame_prefix, args=(''))\n",
    "test_fail(_parse_frame_prefix, args=('subtask if var1==1'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_frbrowse(self, code, kernel, cell):\n",
    "    \"\"\"Display frame interactively.\"\"\"\n",
    "    try:\n",
    "        framename, main_code = _parse_frame_prefix(code)\n",
    "        with Frame(framename):\n",
    "            self.magic_browse(main_code, kernel, cell)\n",
    "    except Exception as e:\n",
    "        print_kernel(f\"frbrowse failed.\\r\\n{e}\", kernel)\n",
    "    return ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[https://github.com/bloomberg/ipydatagrid](https://github.com/bloomberg/ipydatagrid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Head/tail magics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Modeled after [pandas](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.head.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _get_html_data(df):\n",
    "    html = df.convert_dtypes().to_html(notebook=True)\n",
    "    return {'text/html': html}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def _headtail_html(self, df, kernel):\n",
    "    content = {\n",
    "        'data': _get_html_data(df),\n",
    "        'metadata': {},\n",
    "    }\n",
    "    kernel.send_response(kernel.iopub_socket, 'display_data', content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def _magic_headtail(self, code, kernel, cell, tail=False):\n",
    "    try:\n",
    "        expanded_code = macro_expand(code)\n",
    "        df = browse.headtail_get_df(*browse.headtail_df_params(\n",
    "            expanded_code, obs_count(), kernel.nbstata_config.env['missing'], tail=tail\n",
    "        ))\n",
    "        self._headtail_html(df, kernel)\n",
    "    except Exception as e:\n",
    "        print_kernel(f\"{'tail' if tail else 'head'} failed.\\r\\n{e}\", kernel)\n",
    "    return ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_head(self, code, kernel, cell):\n",
    "    \"\"\"Display data in a nicely-formatted table.\"\"\"\n",
    "    return self._magic_headtail(code, kernel, cell, tail=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_frhead(self, code, kernel, cell):\n",
    "    \"\"\"Display data in a nicely-formatted table.\"\"\"\n",
    "    return self._magic_frheadtail(code, kernel, cell, tail=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_tail(self, code, kernel, cell):\n",
    "    \"\"\"Display data in a nicely-formatted table.\"\"\"\n",
    "    return self._magic_headtail(code, kernel, cell, tail=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_frtail(self, code, kernel, cell):\n",
    "    \"\"\"Display data in a nicely-formatted table.\"\"\"\n",
    "    return self._magic_frheadtail(code, kernel, cell, tail=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def _magic_frheadtail(self, code, kernel, cell, tail):\n",
    "    \"\"\"Display frame interactively.\"\"\"\n",
    "    try:\n",
    "        framename, main_code = _parse_frame_prefix(code)\n",
    "        with Frame(framename):\n",
    "            self._magic_headtail(main_code, kernel, cell, tail)\n",
    "    except Exception as e:\n",
    "        print_kernel(f\"{'tail' if tail else 'head'} failed.\\r\\n{e}\", kernel)\n",
    "    return ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Help magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def _get_help_html(self, code):\n",
    "    html_base = \"https://www.stata.com\"\n",
    "    html_help = urllib.parse.urljoin(html_base, \"help.cgi?{}\")\n",
    "    reply = urllib.request.urlopen(html_help.format(code))\n",
    "    html = reply.read().decode(\"utf-8\")\n",
    "\n",
    "    # Remove excessive extra lines (Note css: \"white-space: pre-wrap\")\n",
    "    edited_html = html.replace(\"<p>\\n\", \"<p>\")\n",
    "    soup = bs(edited_html, 'html.parser')\n",
    "\n",
    "    # Set root for links to https://www.stata.com\n",
    "    for a in soup.find_all('a', href=True):\n",
    "        href = a.get('href')\n",
    "        match = re.search(r'{}(.*?)#'.format(code), href)\n",
    "        if match:\n",
    "            hrelative = href.find('#')\n",
    "            a['href'] = href[hrelative:]\n",
    "        elif not href.startswith('http'):\n",
    "            link = a['href']\n",
    "            match = re.search(r'/help.cgi\\?(.+)$', link)\n",
    "            # URL encode bad characters like %\n",
    "            if match:\n",
    "                link = '/help.cgi?'\n",
    "                link += urllib.parse.quote_plus(match.group(1))\n",
    "            a['href'] = urllib.parse.urljoin(html_base, link)\n",
    "            a['target'] = '_blank'\n",
    "\n",
    "    # Remove header 'Stata 15 help for ...'\n",
    "    stata_header = soup.find('h2')\n",
    "    if stata_header:\n",
    "        stata_header.decompose()\n",
    "\n",
    "    # Remove Stata help menu\n",
    "    soup.find('div', id='menu').decompose()\n",
    "\n",
    "    # Remove Copyright notice\n",
    "    copyright = soup.find(string=re.compile(r\".*Copyright.*\", flags=re.DOTALL))\n",
    "    copyright.find_parent(\"table\").decompose()\n",
    "\n",
    "    # Remove last hrule\n",
    "    soup.find_all('hr')[-1].decompose()\n",
    "    \n",
    "    # Remove last br\n",
    "    soup.find_all('br')[-1].decompose()\n",
    "    \n",
    "    # Remove last empty paragraph, empty space\n",
    "    empty_paragraphs = soup.find_all('p', string=\"\")\n",
    "    if str(empty_paragraphs[-1]) == \"<p></p>\":\n",
    "        empty_paragraphs[-1].decompose()\n",
    "\n",
    "    # Set all the backgrounds to transparent\n",
    "    for color in ['#ffffff', '#FFFFFF']:\n",
    "        for bg in ['bgcolor', 'background', 'background-color']:\n",
    "            for tag in soup.find_all(attrs={bg: color}):\n",
    "                if tag.get(bg):\n",
    "                    tag[bg] = 'transparent'\n",
    "\n",
    "    # Set html\n",
    "    css = soup.find('style', {'type': 'text/css'})\n",
    "    with open(self.csshelp_default, 'r') as default:\n",
    "        css.string = default.read()\n",
    "\n",
    "    return str(soup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "@patch_to(StataMagics)\n",
    "def magic_help(self, code, kernel, cell):\n",
    "    \"\"\"Show help file from stata.com/help.cgi?\\{\\}\"\"\"\n",
    "    try:\n",
    "        html_help = self._get_help_html(code)\n",
    "    except Exception as e: # original: (urllib.error.HTTPError, urllib.error.URLError)\n",
    "        msg = \"Failed to fetch HTML help.\\r\\n{0}\"\n",
    "        print_kernel(msg.format(e), kernel)\n",
    "    else:\n",
    "        fallback = 'This front-end cannot display HTML help.'\n",
    "        resp = {\n",
    "            'data': {\n",
    "                'text/html': html_help,\n",
    "                'text/plain': fallback},\n",
    "            'metadata': {}}\n",
    "        kernel.send_response(kernel.iopub_socket, 'display_data', resp)\n",
    "    return ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.display import HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<html>\n",
       "<head>\n",
       "<meta content=\"noindex, nofollow\" name=\"robots\"/>\n",
       "<title>Stata 18 help for order</title>\n",
       "<link href=\"/includes/css/stylehelp.css\" rel=\"stylesheet\" type=\"text/css\"/>\n",
       "<style type=\"text/css\">h2 {\n",
       "    font-family: Arial,Helvetica,Helv,sans-serif;\n",
       "    color: #000000;\n",
       "}\n",
       "\n",
       "pre {\n",
       "    margin: 10px;\n",
       "}\n",
       "\n",
       "table {\n",
       "    background-color: transparent;\n",
       "    border-color: transparent;\n",
       "    bgcolor: transparent;\n",
       "}\n",
       "\n",
       "tr {\n",
       "    background-color: transparent;\n",
       "    border-color: transparent;\n",
       "    bgcolor: transparent;\n",
       "}\n",
       "\n",
       "body {\n",
       "    background-color: transparent;\n",
       "    border-color: transparent;\n",
       "    margin: 0px;\n",
       "}\n",
       "\n",
       "div {\n",
       "    background-color: transparent;\n",
       "    border-color: transparent;\n",
       "}\n",
       "</style>\n",
       "</head>\n",
       "<body>\n",
       "\n",
       "<p><table align=\"center\" bgcolor=\"transparent\" border=\"1\" cellpadding=\"5\" cellspacing=\"0\">\n",
       "<tr bgcolor=\"transparent\"><td align=\"left\" bgcolor=\"transparent\" valign=\"top\">\n",
       "<!-- END HEAD -->\n",
       "\n",
       "<pre>\n",
       "<p><b>[D] order</b> -- Reorder variables in dataset\n",
       "<p><p><a name=\"syntax\"></a><b><u>Syntax</u></b>\n",
       "<p>        <b>order</b> <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> [<b>,</b> <i>options</i>]\n",
       "<p>    <i>options</i>           Description\n",
       "    -------------------------------------------------------------------------\n",
       "    <b>first</b>             move <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> to beginning of dataset; the default\n",
       "    <b>last</b>              move <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> to end of dataset\n",
       "    <b><u>b</u></b><b>efore(</b><a href=\"https://www.stata.com/help.cgi?varname\" target=\"_blank\"><i>varname</i></a><b>)</b>   move <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> before <i>varname</i>\n",
       "    <b><u>a</u></b><b>fter(</b><a href=\"https://www.stata.com/help.cgi?varname\" target=\"_blank\"><i>varname</i></a><b>)</b>    move <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> after <i>varname</i>\n",
       "    <b><u>alpha</u></b><b>betic</b>        alphabetize <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> and move it to beginning of dataset\n",
       "    <b><u>seq</u></b><b>uential</b>        alphabetize <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> keeping numbers sequential and move\n",
       "                        it to beginning of dataset\n",
       "    -------------------------------------------------------------------------\n",
       "<p><p><a name=\"menu\"></a><b><u>Menu</u></b>\n",
       "<p>    <b>Data &gt; Data utilities &gt; Change order of variables</b>\n",
       "<p><p><a name=\"description\"></a><b><u>Description</u></b>\n",
       "<p>    <b>order</b> relocates <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> to a position depending on which option you\n",
       "    specify.  If no option is specified, <b>order</b> relocates <i>varlist</i> to the\n",
       "    beginning of the dataset in the order in which the variables are\n",
       "    specified.\n",
       "<p><p><a name=\"options\"></a><b><u>Options</u></b>\n",
       "<p>    <b>first</b> shifts <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> to the beginning of the dataset.  This is the\n",
       "        default.\n",
       "<p>    <b>last</b> shifts <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> to the end of the dataset.\n",
       "<p>    <b>before(</b><a href=\"https://www.stata.com/help.cgi?varname\" target=\"_blank\"><i>varname</i></a><b>)</b> shifts <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> before <i>varname</i>.\n",
       "<p>    <b>after(</b><a href=\"https://www.stata.com/help.cgi?varname\" target=\"_blank\"><i>varname</i></a><b>)</b> shifts <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> after <i>varname</i>.\n",
       "<p>    <b>alphabetic</b> alphabetizes <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a> and moves it to the beginning of the\n",
       "        dataset.  For example, here is a varlist in <b>alphabetic</b> order:  <b>a x7</b>\n",
       "        <b>x70 x8 x80 z</b>.  If combined with another option, <b>alphabetic</b> just\n",
       "        alphabetizes <i>varlist</i>, and the movement of <i>varlist</i> is controlled by\n",
       "        the other option.\n",
       "<p>    <b>sequential</b> alphabetizes <a href=\"https://www.stata.com/help.cgi?varlist\" target=\"_blank\"><i>varlist</i></a>, keeping variables with the same ordered\n",
       "        letters but with differing appended numbers in sequential order.\n",
       "        <i>varlist</i> is moved to the beginning of the dataset.  For example, here\n",
       "        is a varlist in <b>sequential</b> order: <b>a x7 x8 x70 x80 z</b>.\n",
       "<p><p><a name=\"examples\"></a><b><u>Examples</u></b>\n",
       "<p>    Setup\n",
       "        <b>. webuse auto4</b>\n",
       "<p>    Describe the dataset\n",
       "        <b>. describe</b>\n",
       "<p>    Move <b>make</b> and <b>mpg</b> to the beginning of the dataset\n",
       "        <b>. order make mpg</b>\n",
       "<p>    Describe the dataset\n",
       "        <b>. describe</b>\n",
       "<p>    Make <b>length</b> be the last variable in the dataset\n",
       "        <b>. order length, last</b>\n",
       "<p>    Describe the dataset\n",
       "        <b>. describe</b>\n",
       "<p>    Make <b>weight</b> be the third variable in the dataset\n",
       "        <b>. order weight, before(price)</b>\n",
       "<p>    Describe the dataset\n",
       "        <b>. describe</b>\n",
       "<p>    Alphabetize the variables\n",
       "        <b>. order _all, alphabetic</b>\n",
       "<p>    Describe the dataset\n",
       "        <b>. describe</b>\n",
       "</p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></p></pre>\n",
       "<!-- BEGIN FOOT -->\n",
       "</td></tr></table>\n",
       "\n",
       "\n",
       "\n",
       "</p></body>\n",
       "</html>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_instance = StataMagics()\n",
    "out = test_instance._get_help_html(\"order\")\n",
    "HTML(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}