Add a portserver_address kwarg to pick_unused_port (#12)
This lets the caller specify the address of a portserver without having
to modify the environment.
diff --git a/src/portpicker.py b/src/portpicker.py
index aee7c39..c800161 100644
--- a/src/portpicker.py
+++ b/src/portpicker.py
@@ -133,13 +133,18 @@
IsPortFree = is_port_free # legacy API. pylint: disable=invalid-name
-def pick_unused_port(pid=None):
+def pick_unused_port(pid=None, portserver_address=None):
"""A pure python implementation of PickUnusedPort.
Args:
pid: PID to tell the portserver to associate the reservation with. If
- None,
- the current process's PID is used.
+ None, the current process's PID is used.
+ portserver_address: The address (path) of a unix domain socket
+ with which to connect to a portserver. A leading '@'
+ character indicates an address in the "abstract namespace." If
+ None, or no port is returned by the portserver at the provided
+ address, the environment will be checked for a PORTSERVER_ADDRESS
+ variable. If that's not set, no port server will be used.
Returns:
A port number that is unused on both TCP and UDP.
@@ -149,6 +154,10 @@
_owned_ports.add(port)
return port
# Provide access to the portserver on an opt-in basis.
+ if portserver_address:
+ port = get_port_from_port_server(portserver_address, pid=pid)
+ if port:
+ return port
if 'PORTSERVER_ADDRESS' in os.environ:
port = get_port_from_port_server(os.environ['PORTSERVER_ADDRESS'],
pid=pid)
diff --git a/src/tests/portpicker_test.py b/src/tests/portpicker_test.py
index e799d01..b6ac959 100644
--- a/src/tests/portpicker_test.py
+++ b/src/tests/portpicker_test.py
@@ -70,6 +70,26 @@
@unittest.skipIf('PORTSERVER_ADDRESS' not in os.environ,
'no port server to test against')
+ def testPickUnusedCanSuccessfullyUsePortServerAddressKwarg(self):
+
+ with mock.patch.object(portpicker, '_pick_unused_port_without_server'):
+ portpicker._pick_unused_port_without_server.side_effect = (
+ Exception('eek!')
+ )
+
+ # Since _PickUnusedPortWithoutServer() raises an exception, and
+ # we've temporarily removed PORTSERVER_ADDRESS from os.environ, if
+ # we can successfully obtain a port, the portserver must be working.
+ addr = os.environ.pop('PORTSERVER_ADDRESS')
+ try:
+ port = portpicker.pick_unused_port(portserver_address=addr)
+ self.assertTrue(self.IsUnusedTCPPort(port))
+ self.assertTrue(self.IsUnusedUDPPort(port))
+ finally:
+ os.environ['PORTSERVER_ADDRESS'] = addr
+
+ @unittest.skipIf('PORTSERVER_ADDRESS' not in os.environ,
+ 'no port server to test against')
def testGetPortFromPortServer(self):
"""Exercise the get_port_from_port_server() helper function."""
for _ in range(10):
