fixes #210

Author: jph00

fastcore/_nbdev.py

@@ -160,6 +160,9 @@
          "repo_details": "03_xtras.ipynb",
          "run": "03_xtras.ipynb",
          "do_request": "03_xtras.ipynb",
+         "threaded": "03_xtras.ipynb",
+         "start_server": "03_xtras.ipynb",
+         "start_client": "03_xtras.ipynb",
          "sort_by_run": "03_xtras.ipynb",
          "trace": "03_xtras.ipynb",
          "round_multiple": "03_xtras.ipynb",
@@ -172,7 +175,6 @@
          "parallel": "03_xtras.ipynb",
          "run_procs": "03_xtras.ipynb",
          "parallel_gen": "03_xtras.ipynb",
-         "threaded": "03_xtras.ipynb",
          "lenient_issubclass": "04_dispatch.ipynb",
          "sorted_topologically": "04_dispatch.ipynb",
          "TypeDispatch": "04_dispatch.ipynb",

fastcore/xtras.py

@@ -3,9 +3,9 @@
 __all__ = ['dict2obj', 'repr_dict', 'is_listy', 'shufflish', 'mapped', 'IterLen', 'ReindexCollection', 'open_file',
            'save_pickle', 'load_pickle', 'maybe_open', 'image_size', 'bunzip', 'join_path_file', 'urlwrap', 'urlopen',
            'urlread', 'urljson', 'urlcheck', 'urlclean', 'urlsave', 'urlvalid', 'untar_dir', 'repo_details', 'run',
-           'do_request', 'sort_by_run', 'trace', 'round_multiple', 'modified_env', 'ContextManagers', 'str2bool',
-           'set_num_threads', 'ProcessPoolExecutor', 'ThreadPoolExecutor', 'parallel', 'run_procs', 'parallel_gen',
-           'threaded']
+           'do_request', 'threaded', 'start_server', 'start_client', 'sort_by_run', 'trace', 'round_multiple',
+           'modified_env', 'ContextManagers', 'str2bool', 'set_num_threads', 'ProcessPoolExecutor',
+           'ThreadPoolExecutor', 'parallel', 'run_procs', 'parallel_gen']
 
 # Cell
 from .imports import *
@@ -14,7 +14,7 @@
 from functools import wraps
 
 import mimetypes,pickle,random,json,urllib,subprocess,shlex,bz2,gzip,zipfile,tarfile
-import imghdr,struct,socket,distutils.util,urllib.request,tempfile
+import imghdr,struct,socket,distutils.util,urllib.request,tempfile,time
 from contextlib import contextmanager,ExitStack
 from pdb import set_trace
 from urllib.request import Request
@@ -271,6 +271,43 @@ def do_request(url, post=False, headers=None, **data):
             data = None
     return urljson(Request(url, headers=headers, data=data or None))
 
+# Cell
+def threaded(f):
+    "Run `f` in a thread, and returns the thread"
+    @wraps(f)
+    def _f(*args, **kwargs):
+        res = Thread(target=f, args=args, kwargs=kwargs)
+        res.start()
+        return res
+    return _f
+
+# Cell
+def _socket_det(port,host,dgram):
+    if isinstance(port,int): family,addr = socket.AF_INET,(host or socket.gethostname(),port)
+    else: family,addr = socket.AF_UNIX,port
+    return family,addr,(socket.SOCK_STREAM,socket.SOCK_DGRAM)[dgram]
+
+# Cell
+def start_server(port, host=None, dgram=False, n_queue=None):
+    "Create a `socket` server on `port`, with optional `host`, of type `dgram`"
+    listen_args = [n_queue] if n_queue else []
+    family,addr,typ = _socket_det(port,host,dgram)
+    if family==socket.AF_UNIX:
+        if os.path.exists(addr): os.unlink(addr)
+        assert not os.path.exists(addr), f"{addr} in use"
+    s = socket.socket(family, typ)
+    s.bind(addr)
+    s.listen(*listen_args)
+    return s
+
+# Cell
+def start_client(port, host=None, dgram=False):
+    "Create a `socket` client on `port`, with optional `host`, of type `dgram`"
+    family,addr,typ = _socket_det(port,host,dgram)
+    s = socket.socket(family, typ)
+    s.connect(addr)
+    return s
+
 # Cell
 def _is_instance(f, gs):
     tst = [g if type(g) in [type, 'function'] else g.__class__ for g in gs]
@@ -342,7 +379,6 @@ def str2bool(s):
 # Cell
 from multiprocessing import Process, Queue
 import concurrent.futures
-import time
 from multiprocessing import Manager
 
 # Cell
@@ -446,14 +482,4 @@ def parallel_gen(cls, items, n_workers=defaults.cpus, **kwargs):
     if progress_bar: items = progress_bar(items, leave=False)
     f=partial(_f_pg, cls(**kwargs), queue)
     done=partial(_done_pg, queue, items)
-    yield from run_procs(f, done, L(batches,idx).zip())
-
-# Cell
-def threaded(f):
-    "Run `f` in a thread, and returns the thread"
-    @wraps(f)
-    def _f(*args, **kwargs):
-        res = Thread(target=f, args=args, kwargs=kwargs)
-        res.start()
-        return res
-    return _f
+    yield from run_procs(f, done, L(batches,idx).zip())

nbs/03_xtras.ipynb

@@ -22,7 +22,7 @@
     "from functools import wraps\n",
     "\n",
     "import mimetypes,pickle,random,json,urllib,subprocess,shlex,bz2,gzip,zipfile,tarfile\n",
-    "import imghdr,struct,socket,distutils.util,urllib.request,tempfile\n",
+    "import imghdr,struct,socket,distutils.util,urllib.request,tempfile,time\n",
     "from contextlib import contextmanager,ExitStack\n",
     "from pdb import set_trace\n",
     "from urllib.request import Request\n",
@@ -588,7 +588,7 @@
     {
      "data": {
       "text/plain": [
-       "['e', 'c', 'h', 'b', 'g', 'a', 'd', 'f']"
+       "['a', 'c', 'f', 'd', 'g', 'e', 'b', 'h']"
       ]
      },
      "execution_count": null,
@@ -874,7 +874,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Utilities (other than extensions to Pathlib.Path) for dealing with IO."
+    "Utilities (other than extensions to Pathlib.Path) for dealing with IO and sockets."
    ]
   },
   {
@@ -1411,6 +1411,110 @@
     "    return urljson(Request(url, headers=headers, data=data or None))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "def threaded(f):\n",
+    "    \"Run `f` in a thread, and returns the thread\"\n",
+    "    @wraps(f)\n",
+    "    def _f(*args, **kwargs):\n",
+    "        res = Thread(target=f, args=args, kwargs=kwargs)\n",
+    "        res.start()\n",
+    "        return res\n",
+    "    return _f"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "See `start_client` for an example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "def _socket_det(port,host,dgram):\n",
+    "    if isinstance(port,int): family,addr = socket.AF_INET,(host or socket.gethostname(),port)\n",
+    "    else: family,addr = socket.AF_UNIX,port\n",
+    "    return family,addr,(socket.SOCK_STREAM,socket.SOCK_DGRAM)[dgram]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "def start_server(port, host=None, dgram=False, n_queue=None):\n",
+    "    \"Create a `socket` server on `port`, with optional `host`, of type `dgram`\"\n",
+    "    listen_args = [n_queue] if n_queue else []\n",
+    "    family,addr,typ = _socket_det(port,host,dgram)\n",
+    "    if family==socket.AF_UNIX:\n",
+    "        if os.path.exists(addr): os.unlink(addr)\n",
+    "        assert not os.path.exists(addr), f\"{addr} in use\"\n",
+    "    s = socket.socket(family, typ)\n",
+    "    s.bind(addr)\n",
+    "    s.listen(*listen_args)\n",
+    "    return s"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You can create a TCP client and server pass an int as `port` and optional `host`. `host` defaults to your main network interface if not provided. You can create a Unix socket client and server by passing a string to `port`. A `SOCK_STREAM` socket is created by default, unless you pass `dgram=True`, in which case a `SOCK_DGRAM` socket is created. `n_queue` sets the listening queue size."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#export\n",
+    "def start_client(port, host=None, dgram=False):\n",
+    "    \"Create a `socket` client on `port`, with optional `host`, of type `dgram`\"\n",
+    "    family,addr,typ = _socket_det(port,host,dgram)\n",
+    "    s = socket.socket(family, typ)\n",
+    "    s.connect(addr)\n",
+    "    return s"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "b'hello'\n"
+     ]
+    }
+   ],
+   "source": [
+    "@threaded\n",
+    "def _f():\n",
+    "    with start_server(47354, 'localhost') as s:\n",
+    "        conn,addr = s.accept()\n",
+    "        print(conn.recv(1024))\n",
+    "\n",
+    "_f()\n",
+    "time.sleep(0.2) # Wait for server to start\n",
+    "with start_client(47354, 'localhost') as c: c.send(b\"hello\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -1664,7 +1768,6 @@
     "#export\n",
     "from multiprocessing import Process, Queue\n",
     "import concurrent.futures\n",
-    "import time\n",
     "from multiprocessing import Manager"
    ]
   },
@@ -2045,23 +2148,6 @@
     "test_eq(res.sorted().itemgot(1), x+1)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "#export\n",
-    "def threaded(f):\n",
-    "    \"Run `f` in a thread, and returns the thread\"\n",
-    "    @wraps(f)\n",
-    "    def _f(*args, **kwargs):\n",
-    "        res = Thread(target=f, args=args, kwargs=kwargs)\n",
-    "        res.start()\n",
-    "        return res\n",
-    "    return _f"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},