add retain_types from fastcore.dispatch to fasttransform.utils

Author: Rens Dimmendaal

fasttransform/_modidx.py

@@ -100,4 +100,5 @@
                                      'fasttransform.utils.get_name': ('utils.html#get_name', 'fasttransform/utils.py'),
                                      'fasttransform.utils.is_tuple': ('utils.html#is_tuple', 'fasttransform/utils.py'),
                                      'fasttransform.utils.retain_meta': ('utils.html#retain_meta', 'fasttransform/utils.py'),
-                                     'fasttransform.utils.retain_type': ('utils.html#retain_type', 'fasttransform/utils.py')}}}
+                                     'fasttransform.utils.retain_type': ('utils.html#retain_type', 'fasttransform/utils.py'),
+                                     'fasttransform.utils.retain_types': ('utils.html#retain_types', 'fasttransform/utils.py')}}}

fasttransform/utils.py

@@ -1,7 +1,7 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_utils.ipynb.
 
 # %% auto 0
-__all__ = ['get_name', 'is_tuple', 'retain_meta', 'default_set_meta', 'cast', 'retain_type']
+__all__ = ['get_name', 'is_tuple', 'retain_meta', 'default_set_meta', 'cast', 'retain_type', 'retain_types']
 
 # %% ../nbs/00_utils.ipynb 1
 from typing import Any
@@ -60,3 +60,18 @@ def retain_type(new, old, ret_type=Any,as_copy=False):
         ret_type = old if isinstance(old,type) else type(old)
     if ret_type is NoneType or isinstance(new,ret_type): return new
     return retain_meta(old, cast(new, ret_type), as_copy=as_copy)   
+
+# %% ../nbs/00_utils.ipynb 40
+def retain_types(new, old=None, typs=None):
+    "Cast each item of `new` to type of matching item in `old` if it's a superclass"
+    if not is_listy(new): 
+        typs = Any if typs is None else typs  # make fasttransform.utils.retain_type compatible
+        return retain_type(new, old,typs)
+    if typs is not None:
+        if isinstance(typs, dict):
+            t = first(typs.keys())
+            typs = typs[t]
+        else: t,typs = typs,None
+    else: t = type(old) if old is not None and isinstance(old,type(new)) else type(new)
+    return t(L(new, old, typs).map_zip(retain_types, cycled=True))
+     

nbs/00_utils.ipynb

@@ -432,6 +432,63 @@
     "test_eq(retain_type(FS(1.), None, Any), FS(1.))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "fe4ee917",
+   "metadata": {},
+   "source": [
+    "## Retain types"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "17849108",
+   "metadata": {},
+   "source": [
+    "Copied from fastcore.dispatch, Used in fastai."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f3cd995c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#|export\n",
+    "def retain_types(new, old=None, typs=None):\n",
+    "    \"Cast each item of `new` to type of matching item in `old` if it's a superclass\"\n",
+    "    if not is_listy(new): \n",
+    "        typs = Any if typs is None else typs  # make fasttransform.utils.retain_type compatible\n",
+    "        return retain_type(new, old,typs)\n",
+    "    if typs is not None:\n",
+    "        if isinstance(typs, dict):\n",
+    "            t = first(typs.keys())\n",
+    "            typs = typs[t]\n",
+    "        else: t,typs = typs,None\n",
+    "    else: t = type(old) if old is not None and isinstance(old,type(new)) else type(new)\n",
+    "    return t(L(new, old, typs).map_zip(retain_types, cycled=True))\n",
+    "     "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "37c1d02c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class T(tuple): pass\n",
+    "\n",
+    "t1,t2 = retain_types((1,(1,(1,1))), (2,T((2,T((3,4))))))\n",
+    "test_eq_type(t1, 1)\n",
+    "test_eq_type(t2, T((1,T((1,1)))))\n",
+    "\n",
+    "t1,t2 = retain_types((1,(1,(1,1))), typs = {tuple: [int, {T: [int, {T: [int,int]}]}]})\n",
+    "test_eq_type(t1, 1)\n",
+    "test_eq_type(t2, T((1,T((1,1)))))"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "b712c700",