Switched over from strings for options to using Enums to try and enforce correct options better

Author: CosmicHorrorDev

README.md

@@ -27,22 +27,24 @@ Without an API key pubproxy limits users to one request per second so a `ProxyFe
 ## Quickstart Example
 
 ```python
-from pubproxpy import ProxyFetcher
+from pubproxpy import Level, Protocol, ProxyFetcher
 
 # ProxyFetcher for proxies that use the socks5 protocol, are located in
 # the US or Canada and support POST requests
-socks_pf = ProxyFetcher(protocol="socks5", countries=["US", "CA"], post=True)
+socks_pf = ProxyFetcher(
+    protocol=Protocol.SOCKS5, countries=["US", "CA"], post=True
+)
 
 # ProxyFetcher for proxies that support https, are elite anonymity level,
 # and connected in 15 seconds or less
 https_pf = ProxyFetcher(
-    protocol="http", https=True, level="elite", time_to_connect=15
+    protocol=Protocol.HTTP, https=True, level=Level.ELITE, time_to_connect=15
 )
 
 # Get one socks proxy, followed by 10 https proxies
 # NOTE: even though there are multiple `ProxyFetcher`s the delays are
 #       coordinated between them to prevent rate limiting
-socks_proxy = socks_pf.get()[0]   # Get a single socks proxy
+socks_proxy = socks_pf.get()[0]  # Get a single socks proxy
 https_proxies = https_pf.get(10)  # Get a 10 https proxies
 
 # And then if you want to get any remaining proxies left over before you're
@@ -64,8 +66,8 @@ Since the API doesn't check pretty much anything for correctness, we do our best
 |:--|:--|:--|
 |`exclude_used`|`bool` |[_Default: `True`_] If the `ProxyFetcher` should prevent re-returning proxies|
 |`api_key`|`str`|API key for a paid account, you can also set `$PUBPROXY_API_KEY` to pass your key, passing the `api_key` parameter will override the env-var if both are present|
-|`level`|`str`|[_Options: `"anonymous"`, `"elite"`_] Proxy anonymity level|
-|`protocol`|`str`|[_Options: `"http"`, `"socks4"`, `"socks5"`_] Desired communication protocol|
+|`level`|`pubproxpy.Level`|[_Options: `ANONYMOUS`, `ELITE`_] Proxy anonymity level|
+|`protocol`|`pubproxpy.Protocol`|[_Options: `HTTP`, `SOCKS4`, `SOCKS5`_] Desired communication protocol|
 |`countries`|`str` or `list<str>`|Locations of the proxy using the [ISO-3166 alpha-2](https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2) country code, **Incompatible with `not_countries`**|
 |`not_countries`|`str` or `list<str>`|Blacklist locations of the proxy using the [ISO-3166 alpha-2](https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2) country code, **Incompatible with `countries`**|
 |`last_checked`|`int`|[_Bounds: 1-1000_] Minutes since the proxy was checked|

pubproxpy/__init__.py

@@ -1,5 +1,6 @@
 from pubproxpy.fetcher import ProxyFetcher
+from pubproxpy.types import Level, Protocol
 
 
-__all__ = ["ProxyFetcher"]
+__all__ = ["ProxyFetcher", "Level", "Protocol"]
 __version__ = "1.1.3"

pubproxpy/fetcher.py

@@ -1,8 +1,3 @@
-#!/usr/bin/env python3
-# TODO: check for error with using a proxy to get proxies
-# TODO: `rg` to find any other todos before release
-# TODO: strongly type parameters
-
 import requests
 
 from datetime import datetime as dt
@@ -13,6 +8,7 @@
 
 from pubproxpy._singleton import Singleton
 from pubproxpy.errors import ProxyError, API_ERROR_MAP
+from pubproxpy.types import Level, Protocol
 
 
 class _FetcherShared(metaclass=Singleton):
@@ -30,9 +26,6 @@ def reset(self):
         self.__init__()
 
 
-# TODO: set up tests for things
-# TODO: move all the constants for ProxyFetcher outside of the class?
-#       to constants file maybe?
 class ProxyFetcher:
     """Class used to fetch proxies from the pubproxy API matching the provided
     parameters
@@ -58,12 +51,6 @@ class ProxyFetcher:
         "user_agent",
     )
 
-    # Parameters that have explicit options
-    _PARAM_OPTS = {
-        "level": ("anonymous", "elite"),
-        "protocol": ("http", "socks4", "socks5"),
-    }
-
     # Parameters that are bounded
     _PARAM_BOUNDS = {"last_checked": (1, 1000), "time_to_connect": (1, 60)}
 
@@ -112,6 +99,16 @@ def _verify_params(self, params):
                 " mutually exclusive"
             )
 
+        # Check that protocol and level are the correct type
+        for key, enum_type in (("protocol", Protocol), ("level", Level)):
+            if key in params:
+                val = params[key]
+                if not isinstance(val, enum_type):
+                    raise ValueError(
+                        f"{key} should be of type `{enum_type}` not "
+                        f" `{type(val)}`"
+                    )
+
         # Verify all params are valid, and satisfy the valid bounds or options
         for param, val in params.items():
             if param not in self._PARAMS:
@@ -120,14 +117,6 @@ def _verify_params(self, params):
                     f" {[p for p in self._PARAMS]}"
                 )
 
-            if param in self._PARAM_OPTS:
-                opts = self._PARAM_OPTS[param]
-                if val not in opts:
-                    raise ValueError(
-                        f'invalid value "{val}" for "{param}" options are'
-                        f" {opts}"
-                    )
-
             if param in self._PARAM_BOUNDS:
                 low, high = self._PARAM_BOUNDS[param]
                 if val < low or val > high:
@@ -177,6 +166,11 @@ def _format_params(self, params):
             if isinstance(params["not_country"], (list, tuple)):
                 params["not_country"] = ",".join(params["not_country"])
 
+        # Get value from enums
+        for key in ("level", "type"):
+            if key in params:
+                params[key] = params[key].value
+
         return params
 
     def drain(self):

pubproxpy/types.py

@@ -0,0 +1,12 @@
+from enum import Enum
+
+
+class Level(Enum):
+    ANONYMOUS = "anonymous"
+    ELITE = "elite"
+
+
+class Protocol(Enum):
+    HTTP = "http"
+    SOCKS4 = "socks4"
+    SOCKS5 = "socks5"

tests/test_fetcher.py

@@ -6,7 +6,7 @@
 import os
 from unittest.mock import patch
 
-from pubproxpy import ProxyFetcher
+from pubproxpy import Level, Protocol, ProxyFetcher
 from pubproxpy.fetcher import _FetcherShared
 
 
@@ -94,11 +94,15 @@ def test_params():
     with pytest.raises(ValueError):
         _ = ProxyFetcher(countries="US", not_countries=["CA", "NK"])
 
+    # Switched from strings to `Enum`s when possible
+    with pytest.raises(ValueError):
+        _ = ProxyFetcher(protocol="http")
+
     # And now it's time to check everything
     before_params = {
         "api_key": "<other key>",
-        "level": "elite",
-        "protocol": "http",
+        "level": Level.ELITE,
+        "protocol": Protocol.HTTP,
         "countries": "CA",
         "last_checked": 1,
         "port": 1234,