Merge pull request #33 from ipdata/claude/fix-api-key-env-var-j6bTt

Make pytricia optional and add select_field parameter to lookup

Author: jonathan-kosgei

.github/workflows/python-publish.yml

@@ -1,6 +1,11 @@
 name: Test and Publish ipdata to PyPI
 
-on: push
+on:
+  push:
+    branches:
+      - master
+    tags:
+      - '*'
 
 jobs:
   build-n-publish:

README.md

@@ -57,10 +57,10 @@ Install the latest version of the cli with `pip`.
 pip install ipdata
 ```
 
-or `easy_install`
+To use the [IPTrie](#iptrie) data structure, install with the `trie` extra (requires a C compiler):
 
 ```bash
-easy_install ipdata
+pip install ipdata[trie]
 ```
 
 ## Library Usage
@@ -580,6 +580,8 @@ IPTrie is a production-ready, type-safe trie for IP addresses and CIDR prefixes
 
 ### Quick Start
 
+> **Note:** IPTrie requires the `trie` extra: `pip install ipdata[trie]`
+
 ```python
 from ipdata import IPTrie
 
@@ -741,6 +743,12 @@ A list of possible errors is available at [Status Codes](https://docs.ipdata.co/
 
 ## Tests
 
+Install test dependencies:
+
+```shell
+pip install ipdata[test,trie]
+```
+
 To run all tests
 
 ```shell

setup.cfg

@@ -29,9 +29,13 @@ install_requires =
     click
     click_default_group
     pyperclip
+
+[options.extras_require]
+test =
     pytest
     hypothesis
     python-dotenv
+trie =
     pytricia
 
 [options.entry_points]

src/ipdata/__init__.py

@@ -7,16 +7,20 @@
         >>> ipdata.lookup() # or ipdata.lookup("8.8.8.8") 
 """
 from .ipdata import IPData
-from .iptrie import IPTrie
+
+try:
+    from .iptrie import IPTrie
+except ImportError:
+    IPTrie = None
 
 # Configuration
 api_key = None
 endpoint = "https://api.ipdata.co/"
 default_client = None
 
 
-def lookup(resource="", fields=[]):
-    return _proxy("lookup", resource=resource, fields=fields)
+def lookup(resource="", fields=[], select_field=None):
+    return _proxy("lookup", resource=resource, fields=fields, select_field=select_field)
 
 
 def bulk(resources, fields=[]):

src/ipdata/cli.py

@@ -46,7 +46,6 @@
     ╰───────────────╯ 
 """
 import csv
-import io
 import json
 import logging
 import os
@@ -66,17 +65,16 @@
 from rich.progress import Progress
 from rich.tree import Tree
 
-from .lolcat import LolCat
 from .geofeeds import Geofeed, GeofeedValidationError
-from .ipdata import DotDict, IPData
+from .ipdata import IPData
 
 console = Console()
 
 FORMAT = "%(message)s"
 logging.basicConfig(
     level="ERROR", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()]
 )
-log = logging.getLogger("rich")
+log = logging.getLogger(__name__)
 
 API_KEY_FILE = f"{Path.home()}/.ipdata"
 
@@ -95,19 +93,24 @@ def _lookup(ipdata, *args, **kwargs):
 
 def print_ascii_logo():
     """
-    Print cool ascii logo with lolcat.
+    Print cool ascii logo with rainbow colors.
     """
-    options = DotDict({"animate": False, "os": 6, "spread": 3.0, "freq": 0.1})
-    logo = """
- _           _       _        
-(_)_ __   __| | __ _| |_ __ _ 
+    from rich.text import Text
+
+    logo = r"""
+ _           _       _
+(_)_ __   __| | __ _| |_ __ _
 | | '_ \ / _` |/ _` | __/ _` |
 | | |_) | (_| | (_| | || (_| |
 |_| .__/ \__,_|\__,_|\__\__,_|
   |_|
-    """
-    lol = LolCat()
-    lol.cat(io.StringIO(logo), options)
+"""
+    rainbow = ["red", "orange1", "yellow", "green", "blue", "dark_violet"]
+    lines = logo.strip("\n").split("\n")
+    for i, line in enumerate(lines):
+        text = Text(line)
+        text.stylize(rainbow[i % len(rainbow)])
+        console.print(text)
 
 
 def pretty_print_data(data):

src/ipdata/geofeeds.py

@@ -9,15 +9,11 @@
 from pathlib import Path
 
 import requests
-from rich.logging import RichHandler
 
 from .codes import COUNTRIES, REGION_CODES
 
-FORMAT = "%(message)s"
-logging.basicConfig(
-    level="ERROR", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()]
-)
-log = logging.getLogger("rich")
+log = logging.getLogger(__name__)
+log.addHandler(logging.NullHandler())
 
 pwd = Path(__file__).parent.resolve()
 

src/ipdata/ipdata.py

@@ -26,12 +26,9 @@
 import functools
 
 from requests.adapters import HTTPAdapter, Retry
-from rich.logging import RichHandler
 
-FORMAT = "%(message)s"
-logging.basicConfig(
-    level="ERROR", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()]
-)
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
 
 
 class IPDataException(Exception):
@@ -68,7 +65,7 @@ class IPData(object):
     :param debug: A boolean used to set the log level. Set to True when debugging.
     """
 
-    log = logging.getLogger("rich")
+    log = logging.getLogger(__name__)
 
     valid_fields = {
         "ip",
@@ -100,13 +97,15 @@ class IPData(object):
 
     def __init__(
         self,
-        api_key=os.environ.get("IPDATA_API_KEY"),
+        api_key=None,
         endpoint="https://api.ipdata.co/",
         timeout=60,
         retry_limit=7,
         retry_backoff_factor=1,
         debug=False,
     ):
+        if api_key is None:
+            api_key = os.environ.get("IPDATA_API_KEY")
         if not api_key:
             raise IPDataException("API Key not set. Set an API key via the 'IPDATA_API_KEY' environment variable or see the docs for other ways to do so.")
         # Request settings
@@ -119,6 +118,13 @@ def __init__(
         # Enable debugging
         if debug:
             self.log.setLevel(logging.DEBUG)
+            if not any(
+                not isinstance(h, logging.NullHandler) for h in self.log.handlers
+            ):
+                from rich.logging import RichHandler
+                handler = RichHandler()
+                handler.setFormatter(logging.Formatter("%(message)s", datefmt="[%X]"))
+                self.log.addHandler(handler)
 
         # Work around renamed argument in urllib3.
         if hasattr(urllib3.util.Retry.DEFAULT, "allowed_methods"):
@@ -167,17 +173,18 @@ def _validate_ip_address(self, ip):
         if request_ip.is_private or request_ip.is_reserved or request_ip.is_multicast:
             raise ValueError(f"{ip} is a reserved IP Address")
 
-    def lookup(self, resource="", fields=[]):
+    def lookup(self, resource="", fields=[], select_field=None):
         """
         Makes a GET request to the IPData API for the specified 'resource' and the given 'fields'.
 
         :param resource: Either an IP address or an ASN prefixed by "AS" eg. "AS15169"
         :param fields: A collection of API fields to be returned
+        :param select_field: A single field name to return (convenience alternative to fields)
 
         :returns: An API response as a DotDict object to allow dot notation access of fields eg. data.ip, data.company.name, data.threat.blocklists[0].name etc
 
         :raises IPDataException: if the API call fails or if there is a failure in decoding the response.
-        :raises ValueError: if 'resource' is not a string
+        :raises ValueError: if 'resource' is not a string, or if both 'select_field' and 'fields' are provided
         """
         if type(resource) is not str:
             raise ValueError(f"{resource} must be of type 'str'")
@@ -189,6 +196,13 @@ def lookup(self, resource="", fields=[]):
         if resource and not resource.startswith("AS"):
             self._validate_ip_address(resource)
 
+        if select_field is not None:
+            if fields:
+                raise ValueError("'select_field' and 'fields' are mutually exclusive. Use one or the other.")
+            if not isinstance(select_field, str):
+                raise ValueError("'select_field' must be of type 'str'.")
+            fields = [select_field]
+
         self._validate_fields(fields)
         query_params = self._query_params | {"fields": ",".join(fields)}
 

src/ipdata/iptrie.py

@@ -20,7 +20,10 @@
 from collections.abc import Iterator
 from typing import TypeVar, Generic
 
-from pytricia import PyTricia
+try:
+    from pytricia import PyTricia
+except ImportError:
+    PyTricia = None
 
 T = TypeVar("T")
 
@@ -94,6 +97,11 @@ class IPTrie(Generic[T]):
 
     def __init__(self) -> None:
         """Initialize an empty IPTrie."""
+        if PyTricia is None:
+            raise ImportError(
+                "pytricia is required for IPTrie but is not installed. "
+                "Install it with: pip install ipdata[trie]"
+            )
         self._ipv4: PyTricia = PyTricia(self.IPV4_BITS)
         self._ipv6: PyTricia = PyTricia(self.IPV6_BITS)