fixes #232

fastcore/_nbdev.py

@@ -107,6 +107,9 @@
          "patch_to": "01_basics.ipynb",
          "patch": "01_basics.ipynb",
          "patch_property": "01_basics.ipynb",
+         "ImportEnum": "01_basics.ipynb",
+         "StrEnum": "01_basics.ipynb",
+         "str_enum": "01_basics.ipynb",
          "Stateful": "01_basics.ipynb",
          "PrettyString": "01_basics.ipynb",
          "even_mults": "01_basics.ipynb",

fastcore/basics.py

@@ -11,7 +11,8 @@
            'range_of', 'renumerate', 'first', 'nested_attr', 'nested_idx', 'val2idx', 'uniqueify', 'num_methods',
            'rnum_methods', 'inum_methods', 'fastuple', 'arg0', 'arg1', 'arg2', 'arg3', 'arg4', 'bind', 'map_ex',
            'compose', 'maps', 'partialler', 'instantiate', 'using_attr', 'Self', 'Self', 'copy_func', 'patch_to',
-           'patch', 'patch_property', 'Stateful', 'PrettyString', 'even_mults', 'num_cpus', 'add_props', 'typed']
+           'patch', 'patch_property', 'ImportEnum', 'StrEnum', 'str_enum', 'Stateful', 'PrettyString', 'even_mults',
+           'num_cpus', 'add_props', 'typed']
 
 # Cell
 from .imports import *
@@ -731,6 +732,28 @@ def patch_property(f):
     cls = next(iter(f.__annotations__.values()))
     return patch_to(cls, as_prop=True)(f)
 
+# Cell
+class ImportEnum(enum.Enum):
+    "An `Enum` that can have its values imported"
+    @classmethod
+    def imports(cls):
+        g = sys._getframe(1).f_locals
+        for o in cls: g[o.name]=o
+
+# Cell
+class StrEnum(str,ImportEnum):
+    "An `ImportEnum` that behaves like a `str`"
+    def __str__(self): return self.name
+    @classmethod
+    def imports(cls):
+        g = sys._getframe(1).f_locals
+        for o in cls: g[o.name]=o
+
+# Cell
+def str_enum(name, *vals):
+    "Simplified creation of `StrEnum` types"
+    return StrEnum(name, {o:o for o in vals})
+
 # Cell
 class Stateful:
     "A base class/mixin for objects that should not serialize all their state"

fastcore/imports.py

@@ -1,4 +1,4 @@
-import sys,os,re,typing,itertools,operator,functools,math,warnings,functools,io
+import sys,os,re,typing,itertools,operator,functools,math,warnings,functools,io,enum
 
 from operator import itemgetter,attrgetter
 from warnings import warn

nbs/01_basics.ipynb

@@ -4013,6 +4013,134 @@
     "## Other Helpers"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "class ImportEnum(enum.Enum):\n",
+    "    \"An `Enum` that can have its values imported\"\n",
+    "    @classmethod\n",
+    "    def imports(cls):\n",
+    "        g = sys._getframe(1).f_locals\n",
+    "        for o in cls: g[o.name]=o"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "<h4 id=\"ImportEnum\" class=\"doc_header\"><code>ImportEnum</code><a href=\"\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "\n",
+       "> <code>Enum</code> = []\n",
+       "\n",
+       "An `Enum` that can have its values imported"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "show_doc(ImportEnum, title_level=4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "class StrEnum(str,ImportEnum):\n",
+    "    \"An `ImportEnum` that behaves like a `str`\"\n",
+    "    def __str__(self): return self.name\n",
+    "    @classmethod\n",
+    "    def imports(cls):\n",
+    "        g = sys._getframe(1).f_locals\n",
+    "        for o in cls: g[o.name]=o"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "<h4 id=\"StrEnum\" class=\"doc_header\"><code>StrEnum</code><a href=\"\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "\n",
+       "> <code>Enum</code> = []\n",
+       "\n",
+       "An `ImportEnum` that behaves like a `str`"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "show_doc(StrEnum, title_level=4)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "def str_enum(name, *vals):\n",
+    "    \"Simplified creation of `StrEnum` types\"\n",
+    "    return StrEnum(name, {o:o for o in vals})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "a A\n"
+     ]
+    }
+   ],
+   "source": [
+    "_T = str_enum('_T', 'a', 'b')\n",
+    "test_eq(f'{_T.a}', 'a')\n",
+    "test_eq(_T.a, 'a')\n",
+    "test_eq(list(_T.__members__), ['a','b'])\n",
+    "print(_T.a, _T.a.upper())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "_T = str_enum('_T', 'foobar', 'goobar')\n",
+    "_T.imports()\n",
+    "test_eq(foobar, 'foobar')\n",
+    "test_eq(goobar, 'goobar')"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,