Skip to content

Commit 836b2e7

Browse files
oschwaldoschwald
committed
Add types to records/models
1 parent bdcfd42 commit 836b2e7

File tree

2 files changed

+97
-10
lines changed

2 files changed

+97
-10
lines changed

‎geoip2/models.py‎

Lines changed: 45 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
# pylint: disable=too-many-instance-attributes,too-few-public-methods
1414
import ipaddress
1515
from abc import ABCMeta
16-
from typing import Any, Dict, List, Optional, Union
16+
from typing import Any, cast, Dict, List, Optional, Union
1717

1818
import geoip2.records
1919
from geoip2.mixins import SimpleEquality
@@ -67,6 +67,13 @@ class Country(SimpleEquality):
6767
6868
"""
6969

70+
continent: geoip2.records.Continent
71+
country: geoip2.records.Country
72+
maxmind: geoip2.records.MaxMind
73+
registered_country: geoip2.records.Country
74+
represented_country: geoip2.records.RepresentedCountry
75+
traits: geoip2.records.Traits
76+
7077
def __init__(
7178
self, raw_response: Dict[str, Any], locales: Optional[List[str]] = None
7279
) -> None:
@@ -171,6 +178,11 @@ class City(Country):
171178
172179
"""
173180

181+
city: geoip2.records.City
182+
location: geoip2.records.Location
183+
postal: geoip2.records.Postal
184+
subdivisions: geoip2.records.Subdivisions
185+
174186
def __init__(
175187
self, raw_response: Dict[str, Any], locales: Optional[List[str]] = None
176188
) -> None:
@@ -316,11 +328,14 @@ class SimpleModel(SimpleEquality):
316328

317329
__metaclass__ = ABCMeta
318330

331+
raw: Dict[str, Union[bool, str, int]]
332+
ip_address: str
333+
319334
def __init__(self, raw: Dict[str, Union[bool, str, int]]) -> None:
320335
self.raw = raw
321336
self._network = None
322337
self._prefix_len = raw.get("prefix_len")
323-
self.ip_address = raw.get("ip_address")
338+
self.ip_address = cast(str, raw.get("ip_address"))
324339

325340
def __repr__(self) -> str:
326341
# pylint: disable=no-member
@@ -404,8 +419,14 @@ class AnonymousIP(SimpleModel):
404419
:type: ipaddress.IPv4Network or ipaddress.IPv6Network
405420
"""
406421

407-
def __init__(self, raw: Dict[str, Union[bool, str, int]]) -> None:
408-
super(AnonymousIP, self).__init__(raw)
422+
is_anonymous: bool
423+
is_anonymous_vpn: bool
424+
is_hosting_provider: bool
425+
is_public_proxy: bool
426+
is_tor_exit_node: bool
427+
428+
def __init__(self, raw: Dict[str, bool]) -> None:
429+
super(AnonymousIP, self).__init__(raw) # type: ignore
409430
self.is_anonymous = raw.get("is_anonymous", False)
410431
self.is_anonymous_vpn = raw.get("is_anonymous_vpn", False)
411432
self.is_hosting_provider = raw.get("is_hosting_provider", False)
@@ -446,11 +467,18 @@ class ASN(SimpleModel):
446467
:type: ipaddress.IPv4Network or ipaddress.IPv6Network
447468
"""
448469

470+
autonomous_system_number: Optional[int]
471+
autonomous_system_organization: Optional[str]
472+
449473
# pylint:disable=too-many-arguments
450474
def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
451475
super(ASN, self).__init__(raw)
452-
self.autonomous_system_number = raw.get("autonomous_system_number")
453-
self.autonomous_system_organization = raw.get("autonomous_system_organization")
476+
self.autonomous_system_number = cast(
477+
Optional[int], raw.get("autonomous_system_number")
478+
)
479+
self.autonomous_system_organization = cast(
480+
Optional[str], raw.get("autonomous_system_organization")
481+
)
454482

455483

456484
class ConnectionType(SimpleModel):
@@ -486,9 +514,11 @@ class ConnectionType(SimpleModel):
486514
:type: ipaddress.IPv4Network or ipaddress.IPv6Network
487515
"""
488516

517+
connection_type: Optional[str]
518+
489519
def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
490520
super(ConnectionType, self).__init__(raw)
491-
self.connection_type = raw.get("connection_type")
521+
self.connection_type = cast(Optional[str], raw.get("connection_type"))
492522

493523

494524
class Domain(SimpleModel):
@@ -518,9 +548,11 @@ class Domain(SimpleModel):
518548
519549
"""
520550

551+
domain: Optional[str]
552+
521553
def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
522554
super(Domain, self).__init__(raw)
523-
self.domain = raw.get("domain")
555+
self.domain = cast(Optional[str], raw.get("domain"))
524556

525557

526558
class ISP(ASN):
@@ -568,8 +600,11 @@ class ISP(ASN):
568600
:type: ipaddress.IPv4Network or ipaddress.IPv6Network
569601
"""
570602

603+
isp: Optional[str]
604+
organization: Optional[str]
605+
571606
# pylint:disable=too-many-arguments
572607
def __init__(self, raw: Dict[str, Union[str, int]]) -> None:
573608
super(ISP, self).__init__(raw)
574-
self.isp = raw.get("isp")
575-
self.organization = raw.get("organization")
609+
self.isp = cast(Optional[str], raw.get("isp"))
610+
self.organization = cast(Optional[str], raw.get("organization"))

‎geoip2/records.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ class PlaceRecord(Record):
3232

3333
__metaclass__ = ABCMeta
3434
names: Dict[str, str]
35+
_locales: List[str]
3536

3637
def __init__(
3738
self,
@@ -91,6 +92,9 @@ class City(PlaceRecord):
9192
9293
"""
9394

95+
confidence: Optional[int]
96+
geoname_id: Optional[int]
97+
9498
def __init__(
9599
self,
96100
locales: Optional[List[str]] = None,
@@ -142,6 +146,9 @@ class Continent(PlaceRecord):
142146
143147
"""
144148

149+
code: Optional[str]
150+
geoname_id: Optional[int]
151+
145152
def __init__(
146153
self,
147154
locales: Optional[List[str]] = None,
@@ -207,6 +214,11 @@ class Country(PlaceRecord):
207214
208215
"""
209216

217+
confidence: Optional[int]
218+
geoname_id: Optional[int]
219+
is_in_european_union: bool
220+
iso_code: Optional[str]
221+
210222
def __init__(
211223
self,
212224
locales: Optional[List[str]] = None,
@@ -286,6 +298,8 @@ class RepresentedCountry(Country):
286298
287299
"""
288300

301+
type: Optional[str]
302+
289303
def __init__(
290304
self,
291305
locales: Optional[List[str]] = None,
@@ -371,6 +385,14 @@ class Location(Record):
371385
372386
"""
373387

388+
average_income: Optional[int]
389+
accuracy_radius: Optional[int]
390+
latitude: Optional[float]
391+
longitude: Optional[float]
392+
metro_code: Optional[int]
393+
population_density: Optional[int]
394+
time_zone: Optional[str]
395+
374396
def __init__(
375397
self,
376398
average_income: Optional[int] = None,
@@ -405,6 +427,8 @@ class MaxMind(Record):
405427
406428
"""
407429

430+
queries_remaining: Optional[int]
431+
408432
def __init__(self, queries_remaining: Optional[int] = None, **_) -> None:
409433
self.queries_remaining = queries_remaining
410434

@@ -437,6 +461,9 @@ class Postal(Record):
437461
438462
"""
439463

464+
code: Optional[str]
465+
confidence: Optional[int]
466+
440467
def __init__(
441468
self, code: Optional[str] = None, confidence: Optional[int] = None, **_
442469
) -> None:
@@ -492,6 +519,10 @@ class Subdivision(PlaceRecord):
492519
493520
"""
494521

522+
confidence: Optional[int]
523+
geoname_id: Optional[int]
524+
iso_code: Optional[str]
525+
495526
def __init__(
496527
self,
497528
locales: Optional[List[str]] = None,
@@ -760,6 +791,27 @@ class Traits(Record):
760791
761792
"""
762793

794+
autonomous_system_number: Optional[int]
795+
autonomous_system_organization: Optional[str]
796+
connection_type: Optional[str]
797+
domain: Optional[str]
798+
is_anonymous: bool
799+
is_anonymous_proxy: bool
800+
is_anonymous_vpn: bool
801+
is_hosting_provider: bool
802+
is_legitimate_proxy: bool
803+
is_public_proxy: bool
804+
is_satellite_provider: bool
805+
is_tor_exit_node: bool
806+
isp: Optional[str]
807+
ip_address: Optional[str]
808+
organization: Optional[str]
809+
static_ip_score: Optional[float]
810+
user_count: Optional[int]
811+
user_type: Optional[str]
812+
_network: Optional[str]
813+
_prefix_len: Optional[int]
814+
763815
def __init__(
764816
self,
765817
autonomous_system_number: Optional[int] = None,

0 commit comments

Comments
 (0)