From 2a5469337ae814d648cb54fcbb35a49c7f3d1abb Mon Sep 17 00:00:00 2001
From: ae <85891169+novitae@users.noreply.github.com>
Date: Sun, 20 Oct 2024 13:29:09 +0200
Subject: [PATCH] typing + repr

---
 src/ua_generator/__init__.py      | 12 +++++++-----
 src/ua_generator/data/__init__.py |  5 +++++
 src/ua_generator/user_agent.py    | 19 +++++++++++++++++--
 3 files changed, 29 insertions(+), 7 deletions(-)

diff --git a/src/ua_generator/__init__.py b/src/ua_generator/__init__.py
index dd5cc7f..877c639 100644
--- a/src/ua_generator/__init__.py
+++ b/src/ua_generator/__init__.py
@@ -5,11 +5,13 @@
 """
 import typing
 
-from . import user_agent, options as _options
+from . import user_agent, options as _options, data as _data
 
 
-def generate(device: typing.Union[tuple, str, None] = None,
-             platform: typing.Union[tuple, str, None] = None,
-             browser: typing.Union[tuple, str, None] = None,
-             options: typing.Union[_options.Options, None] = None) -> user_agent.UserAgent:
+def generate(
+    device: typing.Union[tuple[_data._DEVICES_TYPE], _data._DEVICES_TYPE, None] = None,
+    platform: typing.Union[tuple[_data._PLATFORMS_TYPE], _data._PLATFORMS_TYPE, None] = None,
+    browser: typing.Union[tuple[_data._BROWSERS_TYPE], _data._BROWSERS_TYPE, None] = None,
+    options: typing.Union[_options.Options, None] = None
+):
     return user_agent.UserAgent(device, platform, browser, options)
diff --git a/src/ua_generator/data/__init__.py b/src/ua_generator/data/__init__.py
index afca546..da655bb 100644
--- a/src/ua_generator/data/__init__.py
+++ b/src/ua_generator/data/__init__.py
@@ -3,12 +3,17 @@
 Copyright: 2022-2024 Ekin Karadeniz (github.com/iamdual)
 License: Apache License 2.0 
 """
+from typing import Literal
 
 DEVICES = ('desktop', 'mobile')
+_DEVICES_TYPE = Literal['desktop', 'mobile', None]
 
 PLATFORMS = ('windows', 'macos', 'ios', 'linux', 'android')
+_PLATFORMS_TYPE = Literal['windows', 'macos', 'ios', 'linux', 'android', None]
 PLATFORMS_DESKTOP = ('windows', 'macos', 'linux')  # Platforms on desktop devices
 PLATFORMS_MOBILE = ('ios', 'android')  # Platforms on mobile devices
 
 BROWSERS = ('chrome', 'edge', 'firefox', 'safari')
+_BROWSERS_TYPE = Literal['chrome', 'edge', 'firefox', 'safari', None]
 BROWSERS_SUPPORT_CH = ('chrome', 'edge')  # Browsers that support Client Hints
+_BROWSERS_SUPPORT_CH_TYPE = Literal['chrome', 'edge', None]
diff --git a/src/ua_generator/user_agent.py b/src/ua_generator/user_agent.py
index 3a68cd8..ec3049d 100644
--- a/src/ua_generator/user_agent.py
+++ b/src/ua_generator/user_agent.py
@@ -7,14 +7,26 @@
 
 from . import utils, exceptions
 from .client_hints import ClientHints
-from .data import DEVICES, BROWSERS, PLATFORMS, PLATFORMS_DESKTOP, PLATFORMS_MOBILE
+from .data import (
+    DEVICES, _DEVICES_TYPE,
+    BROWSERS, _BROWSERS_TYPE,
+    PLATFORMS, _PLATFORMS_TYPE,
+    PLATFORMS_DESKTOP,
+    PLATFORMS_MOBILE,
+)
 from .data.generator import Generator
 from .headers import Headers
 from .options import Options
 
 
 class UserAgent:
-    def __init__(self, device=None, platform=None, browser=None, options=None):
+    def __init__(
+        self,
+        device: _DEVICES_TYPE = None,
+        platform: _PLATFORMS_TYPE = None,
+        browser: _BROWSERS_TYPE = None,
+        options: typing.Union[Options, None] = None,
+    ):
         self.device: typing.Union[str, None] = utils.choice(device) if device else None
         self.platform: typing.Union[str, None] = utils.choice(platform) if platform else None
         self.browser: typing.Union[str, None] = utils.choice(browser) if browser else None
@@ -88,3 +100,6 @@ def __complete(self):
 
     def __str__(self):
         return self.text
+
+    def __repr__(self) -> str:
+        return f"UserAgent(\"{self.text}\", device={self.device}, platform={self.platform}, browser={self.browser})"
\ No newline at end of file