fixes #155

Author: jph00

fastcore/__init__.py

@@ -1 +1 @@
-__version__ = "1.2.5"
+__version__ = "1.2.6"

fastcore/basics.py

@@ -350,10 +350,14 @@ def range_of(x):
     return list(range(len(x)))
 
 # Cell
-def groupby(x, key):
-    "Like `itertools.groupby` but doesn't need to be sorted, and isn't lazy"
+def groupby(x, key, val=noop):
+    "Like `itertools.groupby` but doesn't need to be sorted, and isn't lazy, plus some extensions"
+    if   isinstance(key,int): key = itemgetter(key)
+    elif isinstance(key,str): key = attrgetter(key)
+    if   isinstance(val,int): val = itemgetter(val)
+    elif isinstance(val,str): val = attrgetter(val)
     res = {}
-    for o in x: res.setdefault(key(o), []).append(o)
+    for o in x: res.setdefault(key(o), []).append(val(o))
     return res
 
 # Cell

nbs/01_basics.ipynb

@@ -222,7 +222,22 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[('hi', ['hi']),\n",
+       " (array(1), [array(1)]),\n",
+       " (1, [1]),\n",
+       " ([1, 2], [1, 2]),\n",
+       " (range(0, 3), [0, 1, 2])]"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "[(o,listify(o)) for o in\n",
     " ('hi',array(1),1, [1,2], range(3))]"
@@ -1978,10 +1993,14 @@
    "outputs": [],
    "source": [
     "#export\n",
-    "def groupby(x, key):\n",
-    "    \"Like `itertools.groupby` but doesn't need to be sorted, and isn't lazy\"\n",
+    "def groupby(x, key, val=noop):\n",
+    "    \"Like `itertools.groupby` but doesn't need to be sorted, and isn't lazy, plus some extensions\"\n",
+    "    if   isinstance(key,int): key = itemgetter(key)\n",
+    "    elif isinstance(key,str): key = attrgetter(key)\n",
+    "    if   isinstance(val,int): val = itemgetter(val)\n",
+    "    elif isinstance(val,str): val = attrgetter(val)\n",
     "    res = {}\n",
-    "    for o in x: res.setdefault(key(o), []).append(o)\n",
+    "    for o in x: res.setdefault(key(o), []).append(val(o))\n",
     "    return res"
    ]
   },
@@ -1994,6 +2013,34 @@
     "test_eq(groupby('aa ab bb'.split(), itemgetter(0)), {'a':['aa','ab'], 'b':['bb']})"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Here's an example of how to *invert* a grouping, using an `int` as `key` (which uses `itemgetter`; passing a `str` will use `attrgetter`), and using a `val` function:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{1: [0], 3: [0, 2], 7: [0], 5: [3, 7], 8: [4], 4: [5]}"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "d = {0: [1, 3, 7], 2: [3], 3: [5], 4: [8], 5: [4], 7: [5]}\n",
+    "groupby(((o,k) for k,v in d.items() for o in v), 0, 1)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

settings.ini

@@ -7,7 +7,7 @@ author = Jeremy Howard and Sylvain Gugger
 author_email = [email protected]
 copyright = fast.ai
 branch = master
-version = 1.2.5
+version = 1.2.6
 min_python = 3.6
 audience = Developers
 language = English