From f9df857a57a2c5fdcaf183e67a695e2c05931c3a Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Wed, 23 Aug 2017 01:33:39 +0530
Subject: [PATCH] Tidy up API.

---
 docs/index.rst     | 13 +++++--------
 econtext/master.py | 12 ++++++++++++
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/docs/index.rst b/docs/index.rst
index 2e43acb8..d24e4b4e 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -78,22 +78,19 @@ to your network topology**.
 
 .. code::
 
-    bastion_host = router.connect(
-        econtext.ssh.Stream,
+    bastion_host = router.ssh(
         hostname='jump-box.mycorp.com'
     )
 
-    ssh_account = router.proxy_connect(
-        bastion_host,
-        econtext.sudo.Stream,
+    ssh_account = router.sudo(
         username='user_with_magic_ssh_key',
         password='sudo password',
+        via=bastion_host,
     )
 
-    internal_box = router.proxy_connect(
-        ssh_account,
-        econtext.ssh.Stream,
+    internal_box = router.ssh(
         hostname='billing0.internal.mycorp.com'
+        via=ssh_account,
     )
         
     internal_box.call(os.system, './run-nightly-billing.py')
diff --git a/econtext/master.py b/econtext/master.py
index 86f20608..3c4fb35c 100644
--- a/econtext/master.py
+++ b/econtext/master.py
@@ -450,7 +450,19 @@ class Router(econtext.core.Router):
         self.register(context, stream)
         return context
 
+    def sudo(self, **kwargs):
+        import econtext.sudo
+        return self.connect(econtext.sudo.Stream, **kwargs)
+
+    def ssh(self, **kwargs):
+        import econtext.ssh
+        return self.connect(econtext.ssh.Stream, **kwargs)
+
     def connect(self, klass, name=None, **kwargs):
+        via = kwargs.pop('via', None)
+        if via is not None:
+            return self.proxy_connect(via, klass, name=name, **kwargs)
+
         context_id = self.alloc_slave_id()
         return self._connect(context_id, klass, name=name, **kwargs)