Astropy `quantity_input` tests

astropy-test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.0.1'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import astropy\n",
    "import astropy.units as u\n",
    "astropy.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@astropy.units.quantity_input\n",
    "def area_of_square(length: u.m, color=\"red\"):\n",
    "    if color == \"blue\":\n",
    "        print(\"That's a particularly nice square.\")\n",
    "    return length**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$9 \\; \\mathrm{m^{2}}$"
      ],
      "text/plain": [
       "<Quantity 9.0 m2>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "area_of_square(3*u.m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "That's a particularly nice square.\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$9 \\; \\mathrm{m^{2}}$"
      ],
      "text/plain": [
       "<Quantity 9.0 m2>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "area_of_square(3*u.m, \"blue\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "Argument 'length' to function 'area_of_square' has no 'unit' attribute. You may want to pass in an astropy Quantity instead.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m~/.anaconda3/envs/pic_dev/lib/python3.6/site-packages/astropy/units/decorators.py\u001b[0m in \u001b[0;36m_validate_arg_value\u001b[0;34m(param_name, func_name, arg, targets, equivalencies)\u001b[0m\n\u001b[1;32m     52\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 53\u001b[0;31m             is_equivalent = arg.unit.is_equivalent(allowed_unit,\n\u001b[0m\u001b[1;32m     54\u001b[0m                                                    equivalencies=equivalencies)\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'int' object has no attribute 'unit'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-1e628acd32f0>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0marea_of_square\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m3\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# here's the meh part\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/.anaconda3/envs/pic_dev/lib/python3.6/site-packages/astropy/utils/decorators.py\u001b[0m in \u001b[0;36marea_of_square\u001b[0;34m(length, color)\u001b[0m\n\u001b[1;32m    872\u001b[0m             \u001b[0mname\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    873\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 874\u001b[0;31m         \u001b[0mfunc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmake_function_with_signature\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mwrapped_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    875\u001b[0m         func = functools.update_wrapper(func, wrapped, assigned=assigned,\n\u001b[1;32m    876\u001b[0m                                         updated=updated)\n",
      "\u001b[0;32m~/.anaconda3/envs/pic_dev/lib/python3.6/site-packages/astropy/units/decorators.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*func_args, **func_kwargs)\u001b[0m\n\u001b[1;32m    217\u001b[0m                 \u001b[0;31m#   the value of the argument:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    218\u001b[0m                 _validate_arg_value(param.name, wrapped_function.__name__,\n\u001b[0;32m--> 219\u001b[0;31m                                     arg, valid_targets, self.equivalencies)\n\u001b[0m\u001b[1;32m    220\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    221\u001b[0m             \u001b[0;31m# Call the original function with any equivalencies in force.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.anaconda3/envs/pic_dev/lib/python3.6/site-packages/astropy/units/decorators.py\u001b[0m in \u001b[0;36m_validate_arg_value\u001b[0;34m(param_name, func_name, arg, targets, equivalencies)\u001b[0m\n\u001b[1;32m     65\u001b[0m             raise TypeError(\"Argument '{0}' to function '{1}' has {2}. \"\n\u001b[1;32m     66\u001b[0m                   \u001b[0;34m\"You may want to pass in an astropy Quantity instead.\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 67\u001b[0;31m                      .format(param_name, func_name, error_msg))\n\u001b[0m\u001b[1;32m     68\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     69\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: Argument 'length' to function 'area_of_square' has no 'unit' attribute. You may want to pass in an astropy Quantity instead."
     ]
    }
   ],
   "source": [
    "area_of_square(3) # here's the meh part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$9 \\; \\mathrm{m^{2}}$"
      ],
      "text/plain": [
       "<Quantity 9.0 m2>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@astropy.units.quantity_input\n",
    "def area_of_square_without_default(length: u.m, color):\n",
    "    if color == \"blue\":\n",
    "        print(\"That's a particularly nice square.\")\n",
    "    return length**2\n",
    "area_of_square_without_default(3*u.m, \"red\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "That's a particularly nice square.\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$9 \\; \\mathrm{m^{2}}$"
      ],
      "text/plain": [
       "<Quantity 9.0 m2>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "area_of_square_without_default(3*u.m, \"blue\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pic_dev]",
   "language": "python",
   "name": "conda-env-pic_dev-py"
  },
  "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}