-
Notifications
You must be signed in to change notification settings - Fork 63
/
entity.py
108 lines (71 loc) · 2.9 KB
/
entity.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import pandas as pd
from warnings import warn
from google.api_core.exceptions import NotFound
from carto.exceptions import CartoException
from ..clients.bigquery_client import BigQueryClient
from ...auth import Credentials, get_default_credentials
try:
from abc import ABC, abstractmethod
except ImportError:
from abc import ABCMeta, abstractmethod
ABC = ABCMeta('ABC', (object,), {'__slots__': ()})
_WORKING_PROJECT = 'carto-do-customers'
class CatalogEntity(ABC):
id_field = 'id'
entity_repo = None
def __init__(self, data):
self.data = data
@property
def id(self):
return self.data[self.id_field]
@classmethod
def get(cls, id_):
return cls.entity_repo.get_by_id(id_)
@classmethod
def get_all(cls, filters=None):
return cls.entity_repo.get_all(filters)
def to_series(self):
return pd.Series(self.data)
def to_dict(self):
return self.data
def __eq__(self, other):
return self.data == other.data
def __ne__(self, other):
return not self == other
def __str__(self):
return '{classname}({data})'.format(classname=self.__class__.__name__, data=self.data.__str__())
def __repr__(self):
return "<{classname}('{entity_id}')>".format(classname=self.__class__.__name__, entity_id=self._get_print_id())
def _get_print_id(self):
if 'slug' in self.data.keys():
return self.data['slug']
return self.id
def _download(self, credentials=None):
credentials = self._get_credentials(credentials)
user_dataset = credentials.get_do_dataset()
bq_client = _get_bigquery_client(_WORKING_PROJECT, credentials)
project, dataset, table = self.id.split('.')
view = 'view_{}_{}'.format(dataset.replace('-', '_'), table)
try:
file_path = bq_client.download_to_file(_WORKING_PROJECT, user_dataset, view)
except NotFound:
raise CartoException('You have not purchased the dataset `{}` yet'.format(self.id))
warn('Data saved: {}.'.format(file_path))
warn("To read it you can do: `pandas.read_csv('{}')`.".format(file_path))
return file_path
def _get_credentials(self, credentials=None):
_credentials = credentials or get_default_credentials()
if not isinstance(credentials, Credentials):
raise ValueError('`credentials` must be a Credentials class instance')
return _credentials
def _get_bigquery_client(project, credentials):
return BigQueryClient(project, credentials)
def is_slug_value(id_value):
return len(id_value.split('.')) == 1
class CatalogList(list):
def __init__(self, data):
super(CatalogList, self).__init__(data)
def get(self, item_id):
return next(filter(lambda item: item.id == item_id, self), None)
def to_dataframe(self):
return pd.DataFrame([item.data for item in self])