Merge pull request #27 from Finnventor/develop

Full Python implementation

.github/workflows/check-python.yml

@@ -0,0 +1,39 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Test Python interface
+
+on:
+  push:
+    branches: [ "develop" ]
+  pull_request:
+    branches: [ "develop" ]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install flake8 pytest
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    - name: Lint with flake8
+      run: |
+        # stop the build if there are Python syntax errors or undefined names
+        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    - name: Test with pytest
+      run: |
+        pytest

python/port4me/__init__.py

@@ -1,41 +1,68 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+from itertools import islice
 import socket
 from getpass import getuser
 from os import getenv
 
 
+__all__ = ["port4me"]
+
+
 # Source: https://chromium.googlesource.com/chromium/src.git/+/refs/heads/master/net/base/port_util.cc
 # Last updated: 2022-10-24
-unsafe_ports_chrome = getenv("PORT4ME_EXCLUDE_UNSAFE_CHROME")
-if unsafe_ports_chrome:
-    unsafe_ports_chrome = set(map(int, unsafe_ports_chrome.split(',')))
-else:
-    unsafe_ports_chrome = {1,7,9,11,13,15,17,19,20,21,22,23,25,37,42,43,53,69,77,79,87,95,101,102,103,104,109,110,111,113,115,117,119,123,135,137,139,143,161,179,389,427,465,512,513,514,515,526,530,531,532,540,548,554,556,563,587,601,636,989,990,993,995,1719,1720,1723,2049,3659,4045,5060,5061,6000,6566,6665,6666,6667,6668,6669,6697,10080}
+unsafe_ports_chrome = getenv("PORT4ME_EXCLUDE_UNSAFE_CHROME", "1,7,9,11,13,15,17,19,20,21,22,23,25,37,42,43,53,69,77,79,87,95,101,102,103,104,109,110,111,113,115,117,119,123,135,137,139,143,161,179,389,427,465,512,513,514,515,526,530,531,532,540,548,554,556,563,587,601,636,989,990,993,995,1719,1720,1723,2049,3659,4045,5060,5061,6000,6566,6665,6666,6667,6668,6669,6697,10080")
 
 # Source: https://www-archive.mozilla.org/projects/netlib/portbanning#portlist
 # Last updated: 2022-10-24
-unsafe_ports_firefox = getenv("PORT4ME_EXCLUDE_UNSAFE_FIREFOX")
-if unsafe_ports_firefox:
-    unsafe_ports_firefox = set(map(int, unsafe_ports_firefox.split(',')))
-else:
-    unsafe_ports_firefox = {1,7,9,11,13,15,17,19,20,21,22,23,25,37,42,43,53,77,79,87,95,101,102,103,104,109,110,111,113,115,117,119,123,135,139,143,179,389,465,512,513,514,515,526,530,531,532,540,556,563,587,601,636,993,995,2049,4045,6000}
+unsafe_ports_firefox = getenv("PORT4ME_EXCLUDE_UNSAFE_FIREFOX", "1,7,9,11,13,15,17,19,20,21,22,23,25,37,42,43,53,77,79,87,95,101,102,103,104,109,110,111,113,115,117,119,123,135,139,143,179,389,465,512,513,514,515,526,530,531,532,540,556,563,587,601,636,993,995,2049,4045,6000")
 
 
 def uint_hash(s):
     h = 0
     for char in s:
-        h = (31 * h + ord(char)) % 2**32 
+        h = (31 * h + ord(char)) % 2**32
     return h
 
 
 def is_port_free(port):
-    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    return (s.connect_ex(('', port)) != 0)
-
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return (s.connect_ex(("", port)) != 0)
+
+
+def get_env_ports(var_name):
+    """Get an ordered set of ports from the environment variable `var_name` and `var_name`_SITE"""
+    ports_dict = {}  # using a dict filled with None here because there is no OrderedSet
+    names = [var_name, var_name+"_SITE"]
+    if var_name == "PORT4ME_EXCLUDE":
+        names.append(var_name+"_UNSAFE")
+
+    for name in names:
+        if name == "PORT4ME_EXCLUDE_UNSAFE":
+            ports = getenv(name, "{chrome},{firefox}")
+        else:
+            ports = getenv(name, "")
+        try:
+            for port in ports.replace("{chrome}", unsafe_ports_chrome).replace("{firefox}", unsafe_ports_firefox).split(","):
+                if port:
+                    port1, _, port2 = port.partition("-")
+                    if port2:
+                        for i in range(int(port1), int(port2)+1):
+                            ports_dict[i] = None
+                    else:
+                        ports_dict[int(port1)] = None
+        except ValueError:
+            raise ValueError("invalid port in environment variable "+name)
+    return ports_dict.keys()
+
+
+def lcg(seed, a=75, c=74, modulus=65537):
+    """
+    Get the next number in a sequence according to a Linear Congruential Generator algorithm.
 
-def LCG(seed, a=75, c=74, modulus=65537):  # constants from the ZX81's algorithm
+    The default constants are from the ZX81.
+    """
     seed %= modulus
     seed_next = (a*seed + c) % modulus
 
@@ -44,54 +71,104 @@ def LCG(seed, a=75, c=74, modulus=65537):  # constants from the ZX81's algorithm
     # seed = modulus-1. To make sure we handle any parameter setup, we
     # detect this manually, increment the seed, and recalculate.
     if seed_next == seed:
-        return LCG(seed+1, a, c, modulus)
+        return lcg(seed+1, a, c, modulus)
 
     #assert 0 <= seed_next <= modulus
     return seed_next
 
 
-def port4me(tool='', user='', min_port=1024, max_port=65535, chrome_safe=True, firefox_safe=True):
+def port4me_gen_unfiltered(tool="", user="", prepend=None):
+    if prepend is None:
+        prepend = get_env_ports("PORT4ME_PREPEND")
+
+    yield from prepend
+
+    if not user:
+        user = getenv("PORT4ME_USER", getuser())
+    if not tool:
+        tool = getenv("PORT4ME_TOOL", "")
+
+    port = uint_hash((user+","+tool).rstrip(","))
+    while True:
+        port = lcg(port)
+        yield port
+
+
+def port4me_gen(tool="", user="", prepend=None, include=None, exclude=None, min_port=1024, max_port=65535):
+    if include is None:
+        include = get_env_ports("PORT4ME_INCLUDE")
+    if exclude is None:
+        exclude = get_env_ports("PORT4ME_EXCLUDE")
+    for port in port4me_gen_unfiltered(tool, user, prepend):
+        if ((min_port <= port <= max_port)
+                and (not include or port in include)
+                and (not exclude or port not in exclude)):
+            yield port
+
+
+_list = list  # necessary to avoid conflicts with list() and the parameter which is named list
+def port4me(tool="", user="", prepend=None, include=None, exclude=None, skip=None, list=0, test=None, max_tries=65536, must_work=True, min_port=1024, max_port=65535):
     """
     Find a free TCP port using a deterministic sequence of ports based on the current username.
-    
+
     This reduces the chance of different users trying to access the same port,
     without having to use a completely random new port every time.
 
     Parameters
     ----------
     tool : str, optional
-        Specify this to get a different port sequence for different tools
+        Used in the seed when generating port numbers, to get a different port sequence for different tools.
     user : str, optional
-        Defaults to determining the username with getuser().
-    min_port: int, optional
+        Used in the seed when generating port numbers. Defaults to determining the username with getuser().
+    prepend : list, optional
+        A list of ports to try first
+    include : list, optional
+        If specified, skip any ports not in this list
+    exclude : list, optional
+        Skip any ports in this list
+    skip : int, optional
+        Skip this many ports at the beginning (after excluded ports have been skipped)
+    list : int, optional
+        Instead of returning a single port, return a list of this many ports without checking if they are free.
+    test : int, optional
+        If specified, return whether the port `test` is not in use. All other parameters will be ignored.
+    max_tries : int, optional
+        Raise a TimeoutError if it takes more than this many tries to find a port. Default is 65536.
+    must_work : bool, optional
+        If True, then an error is produced if no port could be found. If False, then `-1` is returned.
+    min_port : int, optional
         Skips any ports that are smaller than this
-    max_port: int, optional
+    max_port : int, optional
         Skips any ports that are larger than this
-    chrome_safe: bool, optional
-        Whether to skip ports that Chrome refuses to open
-    firefox_safe: bool, optional
-        Whether to skip ports that Firefox refuses to open
-
-    See Also
-    --------
-    `unsafe_ports_chrome` : set of ports to be skipped when `chrome_safe=True`
-    `unsafe_ports_firefox` : set of ports to be skipped when `firefox_safe=True`
     """
-    if not user: user = getuser()
+    if test:
+        return is_port_free(test)
+
+    tries = 1
+
+    gen = port4me_gen(tool, user, prepend, include, exclude, min_port, max_port)
+
+    if skip is None:
+        skip = getenv("PORT4ME_SKIP", 0)
+    gen = islice(gen, skip, None)
+
+    if list:
+        return _list(islice(gen, list))
+
+    for port in gen:
+        if is_port_free(port):
+            break
 
-    port = uint_hash((user+','+tool).rstrip(','))
+        if max_tries and tries > max_tries:
+            if must_work:
+                raise TimeoutError("Failed to find a free TCP port after {} attempts".format(max_tries))
+            else:
+                return -1
 
-    while (not (min_port <= port <= max_port)
-           or (chrome_safe and port in unsafe_ports_chrome)
-           or (firefox_safe and port in unsafe_ports_firefox)
-           or not is_port_free(port)):
-        port = LCG(port)
+        tries += 1
 
     return port
 
 
-if __name__ == '__main__':
-    print(port4me(user='alice'))
-    print(port4me('rstudio', user='alice'))
-    print(port4me('jupyter-notebook', user='alice'))  # gets incorrect result
-    print(port4me(user='bob'))
+if __name__ == "__main__":
+    print(port4me())

python/pyproject.toml

@@ -20,3 +20,8 @@ classifiers = [
 [project.urls]
 "Homepage" = "https://github.com/HenrikBengtsson/port4me"
 "Bug Tracker" = "https://github.com/HenrikBengtsson/port4me/issues"
+
+[tool.pytest.ini_options]
+addopts = [
+    "--import-mode=importlib",
+]

python/tests/__init__.py

@@ -0,0 +1 @@
+

python/tests/test_port4me.py

@@ -0,0 +1,21 @@
+from port4me import port4me
+
+
+def test_alice():
+    assert port4me('', 'alice', list=5) == [30845, 19654, 32310, 63992, 15273]
+
+
+def test_alice_tool():
+    assert port4me('jupyter-notebook', 'alice', list=1) == [29525]
+
+
+def test_alice_prepend():
+    assert port4me('', 'alice', list=5, prepend=[9876, 5432]) == [9876, 5432, 30845, 19654, 32310]
+
+
+def test_bob():
+    assert port4me('', 'bob', list=5) == [54242, 4930, 42139, 14723, 55707]
+
+
+def test_bob_skip_exclude():
+    assert port4me('', 'bob', list=2, skip=1, exclude=[54242, 14723]) == [42139, 55707]

python/tests/test_uint_hash.py

@@ -0,0 +1,13 @@
+from port4me import uint_hash
+
+def test_empty():
+    assert uint_hash('') == 0
+
+def test_A():
+    assert uint_hash('A') == 65
+
+def test_alice():
+    assert uint_hash('alice,rstudio') == 3688618396
+
+def test_long():
+    assert uint_hash('port4me - get the same, personal, free TCP port over and over') == 1731535982