From 462cd57407a800268c791f82123ad1c11892e967 Mon Sep 17 00:00:00 2001 From: zaanposni Date: Mon, 16 Dec 2019 15:31:09 +0100 Subject: [PATCH] added option to get response obj of api request returned. fixes #32 --- vvspy/__init__.py | 43 ++++++++++++++++++++++++++++++------------- vvspy/arrivals.py | 9 ++++++++- vvspy/departures.py | 9 ++++++++- vvspy/trip.py | 8 +++++++- 4 files changed, 53 insertions(+), 16 deletions(-) diff --git a/vvspy/__init__.py b/vvspy/__init__.py index c4e27a9..20202fb 100644 --- a/vvspy/__init__.py +++ b/vvspy/__init__.py @@ -1,6 +1,7 @@ from datetime import datetime as __datetime from typing import List as __List from typing import Union as __Union +from requests.models import Response as __Response from .obj import Arrival as __Arrival from .obj import Departure as __Departure @@ -10,8 +11,8 @@ from .arrivals import get_arrivals -def departures_now(station_id: __Union[str, int], limit: int = 100, - **kwargs) -> __List[__Union[__Arrival, __Departure]]: +def departures_now(station_id: __Union[str, int], limit: int = 100, return_resp: bool = False, + **kwargs) -> __Union[__List[__Departure], __Response, None]: """ Same as `get_departures` But `datetime.datetime.now()` is already used as parameter. @@ -20,11 +21,13 @@ def departures_now(station_id: __Union[str, int], limit: int = 100, Returns none on webrequest errors or no results found. """ - return get_departures(station_id=station_id, check_time=__datetime.now(), limit=limit, **kwargs) + return get_departures(station_id=station_id, check_time=__datetime.now(), limit=limit, + return_resp=return_resp, **kwargs) def get_departure(station_id: __Union[str, int], check_time: __datetime = None, debug: bool = False, - request_params: dict = None, **kwargs) -> __Union[__Departure, None]: + request_params: dict = None, return_resp: bool = False, **kwargs)\ + -> __Union[__Departure, __Response, None]: """ Same as `get_departures` But limited to one obj as result. @@ -34,8 +37,12 @@ def get_departure(station_id: __Union[str, int], check_time: __datetime = None, """ try: - return get_departures(station_id=station_id, check_time=check_time, limit=1, debug=debug, - request_params=request_params, **kwargs)[0] + if return_resp: + return get_departures(station_id=station_id, check_time=check_time, limit=1, debug=debug, + request_params=request_params, return_resp=return_resp, **kwargs) + else: + return get_departures(station_id=station_id, check_time=check_time, limit=1, debug=debug, + request_params=request_params, return_resp=return_resp, **kwargs)[0] except IndexError: # no results returned if debug: print("No departures found.") @@ -47,7 +54,8 @@ def get_departure(station_id: __Union[str, int], check_time: __datetime = None, def get_arrival(station_id: __Union[str, int], check_time: __datetime = None, debug: bool = False, - request_params: dict = None, **kwargs) -> __Union[__Arrival, None]: + request_params: dict = None, return_resp: bool = False, **kwargs)\ + -> __Union[__Arrival, __Response, None]: """ Same as `get_arrivals` But limited to one obj as result. @@ -57,8 +65,12 @@ def get_arrival(station_id: __Union[str, int], check_time: __datetime = None, de """ try: - return get_arrivals(station_id=station_id, check_time=check_time, limit=1, debug=debug, - request_params=request_params, **kwargs)[0] + if return_resp: + return get_arrivals(station_id=station_id, check_time=check_time, limit=1, debug=debug, + request_params=request_params, return_resp=return_resp, **kwargs) + else: + return get_arrivals(station_id=station_id, check_time=check_time, limit=1, debug=debug, + request_params=request_params, return_resp=return_resp, **kwargs)[0] except IndexError: # no results returned if debug: print("No arrivals found.") @@ -70,8 +82,8 @@ def get_arrival(station_id: __Union[str, int], check_time: __datetime = None, de def get_trip(origin_station_id: __Union[str, int], destination_station_id: __Union[str, int], - check_time: __datetime = None, debug: bool = False, - request_params: dict = None, **kwargs) -> __Union[__Trip, None]: + check_time: __datetime = None, debug: bool = False, request_params: dict = None, + return_resp: bool = False, **kwargs) -> __Union[__Trip, __Response, None]: """ Same as `get_trips` But limited to one obj as result. @@ -81,8 +93,13 @@ def get_trip(origin_station_id: __Union[str, int], destination_station_id: __Uni """ try: - return get_trips(origin_station_id=origin_station_id, destination_station_id=destination_station_id, - check_time=check_time, limit=1, debug=debug, request_params=request_params, **kwargs)[0] + if return_resp: + return get_trips(origin_station_id=origin_station_id, destination_station_id=destination_station_id, + check_time=check_time, limit=1, debug=debug, request_params=request_params, + return_resp=return_resp, **kwargs) + else: + return get_trips(origin_station_id=origin_station_id, destination_station_id=destination_station_id, + check_time=check_time, limit=1, debug=debug, request_params=request_params, **kwargs)[0] except IndexError: # no results returned if debug: print("No trips found.") diff --git a/vvspy/arrivals.py b/vvspy/arrivals.py index 727df42..4b398ab 100644 --- a/vvspy/arrivals.py +++ b/vvspy/arrivals.py @@ -1,6 +1,7 @@ from typing import List, Union from datetime import datetime import requests +from requests.models import Response import json import traceback @@ -10,7 +11,8 @@ def get_arrivals(station_id: Union[str, int], check_time: datetime = None, limit: int = 100, debug: bool = False, - request_params: dict = None, **kwargs) -> Union[List[Arrival], None]: + request_params: dict = None, return_resp: bool = False, **kwargs)\ + -> Union[List[Arrival], Response, None]: r""" Returns: List[:class:`vvspy.obj.Arrival`] @@ -48,6 +50,8 @@ def get_arrivals(station_id: Union[str, int], check_time: datetime = None, limit request_params Optional[:class:`dict`] params parsed to the api request (e.g. proxies) default {} + return_resp Optional[:class:`bool`] + if set, the function returns the response object of the API request. kwargs Optional[:class:`dict`] Check arrivals.py to see all available kwargs. """ @@ -97,6 +101,9 @@ def get_arrivals(station_id: Union[str, int], check_time: datetime = None, limit print(f"{r.text}") return + if return_resp: + return r + try: r.encoding = 'UTF-8' return _parse_response(r.json()) # TODO: error handling diff --git a/vvspy/departures.py b/vvspy/departures.py index 250f6df..2da95b0 100644 --- a/vvspy/departures.py +++ b/vvspy/departures.py @@ -1,6 +1,7 @@ from typing import List, Union from datetime import datetime import requests +from requests.models import Response import json import traceback @@ -11,7 +12,8 @@ def get_departures(station_id: Union[str, int], check_time: datetime = None, limit: int = 100, debug: bool = False, - request_params: dict = None, **kwargs) -> Union[List[Departure], None]: + request_params: dict = None, return_resp: bool = False, **kwargs)\ + -> Union[List[Departure], Response, None]: r""" Returns: List[:class:`vvspy.obj.Departure`] @@ -49,6 +51,8 @@ def get_departures(station_id: Union[str, int], check_time: datetime = None, lim request_params Optional[:class:`dict`] params parsed to the api request (e.g. proxies) default {} + return_resp Optional[:class:`bool`] + if set, the function returns the response object of the API request. kwargs Optional[:class:`dict`] Check departures.py to see all available kwargs. @@ -98,6 +102,9 @@ def get_departures(station_id: Union[str, int], check_time: datetime = None, lim print(f"{r.text}") return + if return_resp: + return r + try: r.encoding = 'UTF-8' return _parse_response(r.json()) # TODO: error handling diff --git a/vvspy/trip.py b/vvspy/trip.py index 01cb27d..fdc445f 100644 --- a/vvspy/trip.py +++ b/vvspy/trip.py @@ -1,5 +1,6 @@ from datetime import datetime, timezone import requests +from requests.models import Response import json from typing import Union, List import traceback @@ -11,7 +12,7 @@ def get_trips(origin_station_id: Union[str, int], destination_station_id: Union[str, int], check_time: datetime = None, limit: int = 100, debug: bool = False, - request_params: dict = None, **kwargs) -> Union[List[Trip], None]: + request_params: dict = None, return_resp: bool = False, **kwargs) -> Union[List[Trip], Response, None]: r""" Returns: List[:class:`vvspy.obj.Trip`] @@ -49,6 +50,8 @@ def get_trips(origin_station_id: Union[str, int], destination_station_id: Union[ request_params Optional[:class:`dict`] params parsed to the api request (e.g. proxies) default {} + return_resp Optional[:class:`bool`] + if set, the function returns the response object of the API request. kwargs Optional[:class:`dict`] Check trips.py to see all available kwargs. """ @@ -118,6 +121,9 @@ def get_trips(origin_station_id: Union[str, int], destination_station_id: Union[ print(f"{r.text}") return + if return_resp: + return r + try: r.encoding = 'UTF-8' return _parse_response(r.json(), limit=limit) # TODO: error handling