tjorim · tjorim · Jan 1, 2025 · Dec 30, 2024 · Dec 30, 2024 · Dec 30, 2024
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -0,0 +1,14 @@
+{
+    "name": "Python Dev Container",
+    "image": "mcr.microsoft.com/vscode/devcontainers/python:3.9",
+    "features": {},
+    "customizations": {
+        "vscode": {
+            "extensions": [
+                "ms-python.python",
+                "ms-python.vscode-pylance"
+            ]
+        }
+    },
+    "postCreateCommand": "pip install -r requirements.txt"
+}
diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
diff --git a/README.md b/README.md
@@ -1,3 +1,70 @@
 # pyRail
 
-A Python wrapper for the iRail API
+A Python wrapper for the iRail API.
+
+## Overview
+
+pyRail is a Python library that provides a convenient interface for interacting with the iRail API. It supports various endpoints such as stations, liveboard, vehicle, connections, and disturbances. The library also includes features like caching and rate limiting to optimize API usage.
+
+## Installation
+
+To install pyRail, use pip:
+
+```sh
+pip install pyrail
+```
+
+## Usage
+Here is an example of how to use pyRail:
+
+```python
+from pyrail.irail import iRail
+
+# Create an instance of the iRail class
+api = iRail(format='json', lang='en')
+
+# Make a request to the 'stations' endpoint
+response = api.do_request('stations')
+
+# Print the response
+print(response)
+```
+
+## Features
+
+- Supports multiple endpoints: stations, liveboard, vehicle, connections, disturbances
+- Caching and conditional GET requests using ETag
+- Rate limiting to handle API rate limits
+
+## Configuration
+
+You can configure the format and language for the API requests:
+
+```python
+api = iRail(format='json', lang='en')
+```
+
+Supported formats: json, xml, jsonp
+
+Supported languages: nl, fr, en, de
+
+## Logging
+
+You can set the logging level at runtime to get detailed logs:
+
+```python
+api.set_logging_level(logging.DEBUG)
+```
+
+## Contributing
+Contributions are welcome! Please open an issue or submit a pull request.
+
+## Contributors
+- @tjorim
+- @jcoetsie
+
+## License
+
+This project is licensed under the MIT License. See the LICENSE file for details.
+
+
diff --git a/pyrail/irail.py b/pyrail/irail.py
@@ -1,4 +1,10 @@
+import logging
 import requests
+import time
+from threading import Lock
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 session = requests.Session()
 
@@ -16,13 +22,15 @@
 
 class iRail:
 
-    def __init__(self, format=None, lang=None):
-        if format is None:
-            format = 'json'
+    def __init__(self, format='json', lang='en'):
         self.format = format
-        if lang is None:
-            lang = 'en'
    <
8000
span class='blob-code-inner blob-code-marker ' data-code-marker=" ">        self.lang = lang
+        self.tokens = 3
+        self.burst_tokens = 5
+        self.last_request_time = time.time()
+        self.lock = Lock()
+        self.etag_cache = {}
+        logger.info("iRail instance created")
 
     @property
     def format(self):
@@ -46,28 +54,81 @@ def lang(self, value):
         else:
             self.__lang = 'en'
 
+    def _refill_tokens(self):
+        logger.debug("Refilling tokens")
+        current_time = time.time()
+        elapsed = current_time - self.last_request_time
+        self.last_request_time = current_time
+
+        # Refill tokens based on elapsed time
+        self.tokens += elapsed * 3  # 3 tokens per second
+        if self.tokens > 3:
+            self.tokens = 3
+
+        # Refill burst tokens
+        self.burst_tokens += elapsed * 3  # 3 burst tokens per second
+        if self.burst_tokens > 5:
+            self.burst_tokens = 5
+
     def do_request(self, method, args=None):
+        logger.info(f"Starting request to endpoint: {method}")
+        with self.lock:
+            self._refill_tokens()
+
+            if self.tokens < 1:
+                if self.burst_tokens >= 1:
+                    self.burst_tokens -= 1
+                else:
+                    logger.warning("Rate limiting, waiting for tokens")
+                    time.sleep(1 - (time.time() - self.last_request_time))
+                    self._refill_tokens()
+                    self.tokens -= 1
+            else:
+                self.tokens -= 1
+
         if method in methods:
             url = base_url.format(method)
             params = {'format': self.format, 'lang': self.lang}
             if args:
                 params.update(args)
+            headers = {}
+
+            # Add If-None-Match header if we have a cached ETag
+                if method in self.etag_cache:
+                logger.debug(f"Adding If-None-Match header with value: {self.etag_cache[method]}")
+                headers['If-None-Match'] = self.etag_cache[method]
+
             try:
                 response = session.get(url, params=params, headers=headers)
-                try:
-                    json_data = response.json()
-                    return json_data
-                except ValueError:
-                    return -1
+                if response.status_code == 429:
+                    logger.warning("Rate limited, waiting for retry-after header")
+                    retry_after = int(response.headers.get("Retry-After", 1))
+                    time.sleep(retry_after)
+                    return self.do_request(method, args)
+                if response.status_code == 200:
+                    # Cache the ETag from the response
+                    if 'Etag' in response.headers:
+                        self.etag_cache[method] = response.headers['Etag']
+                    try:
+                        json_data = response.json()
+                        return json_data
+                    except ValueError:
+                        return -1
+                elif response.status_code == 304:
+                    logger.info("Data not modified, using cached data")
+                    return None
+                else:
+                    logger.error(f"Request failed with status code: {response.status_code}")
+                    return None
             except requests.exceptions.RequestException as e:
-                print(e)
+                logger.error(f"Request failed: {e}")
                 try:
                     session.get('https://1.1.1.1/', timeout=1)
                 except requests.exceptions.ConnectionError:
-                    print("Your internet connection doesn't seem to be working.")
+                    logger.error("Internet connection failed")
                     return -1
                 else:
-                    print("The iRail API doesn't seem to be working.")
+                    logger.error("iRail API failed")
                     return -1
 
     def get_stations(self):

diff --git a/tests/__init__.py b/tests/__init__.py
@@ -0,0 +1 @@
+from pyrail.irail import irail
diff --git a/tests/test_irail.py b/tests/test_irail.py
@@ -1,7 +1,26 @@
-from pyrail import iRail
+import unittest
+from unittest.mock import patch, MagicMock
+from pyrail.irail import iRail
 
-def test_irail_station():
+class TestiRailAPI(unittest.TestCase):
 
-    irail_instance = iRail()
-    response = irail_instance.get_stations()
-    print(response)
+    @patch('requests.Session.get')
+    def test_successful_request(self, mock_get):
+        # Mock the response to simulate a successful request
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = {'data': 'some_data'}
+        mock_get.return_value = mock_response
+
+        irail_instance = iRail()
+
+        # Call the method that triggers the API request
+        response = irail_instance.do_request('stations')
+
+        # Check that the request was successful
+        self.assertEqual(mock_get.call_count, 1, "Expected one call to the requests.Session.get method")
+        self.assertEqual(response, {'data': 'some_data'}, "Expected response data to match the mocked response")
+
+
+if __name__ == '__main__':
+    unittest.main()
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		from pyrail.irail import irail
Copy link Contributor coderabbitai bot Dec 30, 2024 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion Unused import. According to static analysis, `from pyrail.irail import irail` is unused. Remove or alias it to avoid confusion. -from pyrail.irail import irail 🧰 Tools 🪛 Ruff (0.8.2) 1-1: `pyrail.irail.irail` imported but unused; consider removing, adding to `__all__`, or using a redundant alias (F401)