fixes #746

Author: jph00

fastcore/foundation.py

@@ -92,9 +92,9 @@ def product(xs):
     return reduce(operator.mul, [o for o in xs if o is not None], 1)
 
 # %% ../nbs/02_foundation.ipynb
-def flatmap(f, xs):
+def flatmap(f, xs, **kwargs):
     "Apply f to each element and flatten the results into a single list."
-    return [y for x in xs for y in f(x)]
+    return [y for x in xs for y in f(x, **kwargs)]
 
 # %% ../nbs/02_foundation.ipynb
 class CollBase:
@@ -473,9 +473,9 @@ def setattrs(self:L, attr, val):
 
 # %% ../nbs/02_foundation.ipynb
 @patch
-def flatmap(self:L, f):
+def flatmap(self:L, f, **kwargs):
     "Apply f to each element and flatten the results into a single L."
-    return L(flatmap(f, self))
+    return L(flatmap(f, self, **kwargs))
 
 # %% ../nbs/02_foundation.ipynb
 @patch

nbs/02_foundation.ipynb

@@ -285,12 +285,26 @@
        "\n",
        "### is_iter\n",
        "\n",
-       ">      is_iter (o)\n",
+       "```python\n",
+       "\n",
+       "def is_iter(\n",
+       "    o\n",
+       "):\n",
+       "\n",
+       "\n",
+       "```\n",
        "\n",
        "*Test whether `o` can be used in a `for` loop*"
       ],
       "text/plain": [
-       ">      is_iter (o)\n",
+       "```python\n",
+       "\n",
+       "def is_iter(\n",
+       "    o\n",
+       "):\n",
+       "\n",
+       "\n",
+       "```\n",
        "\n",
        "*Test whether `o` can be used in a `for` loop*"
       ]
@@ -591,9 +605,9 @@
    "outputs": [],
    "source": [
     "#| export\n",
-    "def flatmap(f, xs):\n",
+    "def flatmap(f, xs, **kwargs):\n",
     "    \"Apply f to each element and flatten the results into a single list.\"\n",
-    "    return [y for x in xs for y in f(x)]"
+    "    return [y for x in xs for y in f(x, **kwargs)]"
    ]
   },
   {
@@ -830,6 +844,35 @@
     "flatmap(factpairs, [6,10])"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "f0026d6c",
+   "metadata": {},
+   "source": [
+    "You can also use kwargs, for instance to apply `str.split` with a custom separator:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "754fff88",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['a', 'b', 'c', 'd', 'e']"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "flatmap(str.split, [\"a-b-c\", \"d-e\"], sep=\"-\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "d8a026cf",
@@ -1303,12 +1346,26 @@
        "\n",
        "### L.__getitem__\n",
        "\n",
-       ">      L.__getitem__ (idx)\n",
+       "```python\n",
+       "\n",
+       "def __getitem__(\n",
+       "    idx\n",
+       "):\n",
+       "\n",
+       "\n",
+       "```\n",
        "\n",
        "*Retrieve `idx` (can be list of indices, or mask, or int) items*"
       ],
       "text/plain": [
-       ">      L.__getitem__ (idx)\n",
+       "```python\n",
+       "\n",
+       "def __getitem__(\n",
+       "    idx\n",
+       "):\n",
+       "\n",
+       "\n",
+       "```\n",
        "\n",
        "*Retrieve `idx` (can be list of indices, or mask, or int) items*"
       ]
@@ -1352,12 +1409,26 @@
        "\n",
        "### L.__setitem__\n",
        "\n",
-       ">      L.__setitem__ (idx, o)\n",
+       "```python\n",
+       "\n",
+       "def __setitem__(\n",
+       "    idx, o\n",
+       "):\n",
+       "\n",
+       "\n",
+       "```\n",
        "\n",
        "*Set `idx` (can be list of indices, or mask, or int) items to `o` (which is broadcast if not iterable)*"
       ],
       "text/plain": [
-       ">      L.__setitem__ (idx, o)\n",
+       "```python\n",
+       "\n",
+       "def __setitem__(\n",
+       "    idx, o\n",
+       "):\n",
+       "\n",
+       "\n",
+       "```\n",
        "\n",
        "*Set `idx` (can be list of indices, or mask, or int) items to `o` (which is broadcast if not iterable)*"
       ]
@@ -3089,9 +3160,9 @@
    "source": [
     "#| export\n",
     "@patch\n",
-    "def flatmap(self:L, f):\n",
+    "def flatmap(self:L, f, **kwargs):\n",
     "    \"Apply f to each element and flatten the results into a single L.\"\n",
-    "    return L(flatmap(f, self))"
+    "    return L(flatmap(f, self, **kwargs))"
    ]
   },
   {
@@ -3112,6 +3183,24 @@
     "test_eq(L(\"a,b,c\", \"d,e\").flatmap(Self.split(',')), ['a', 'b', 'c', 'd', 'e'])"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "94aebdb0",
+   "metadata": {},
+   "source": [
+    "or alternatively use kwargs:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "083f30df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_eq(L(\"a-b-c\", \"d-e\").flatmap(str.split, sep='-'), ['a', 'b', 'c', 'd', 'e'])"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "98e3c09e",
@@ -3141,6 +3230,27 @@
     "L(\"a,b,c\", \"d,e\").map(Self.split(',')).concat()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ed763699",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['a', 'b', 'c', 'd', 'e']"
+      ]
+     },
+     "execution_count": null,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "L(\"a-b-c\", \"d-e\").map(str.split, sep='-').concat()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "d0e47eb0",
@@ -4319,13 +4429,7 @@
    ]
   }
  ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "python3",
-   "language": "python",
-   "name": "python3"
-  }
- },
+ "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
 }