Skip to content

Commit 58ed7f2

Browse files
amalek215amalek215
authored
Merge pull request #1 from FlipperPA/connect_defaults
Linting changes.
2 parents f293f07 + 0fcff8a commit 58ed7f2

File tree

2 files changed

+107
-87
lines changed

2 files changed

+107
-87
lines changed

‎wrds/sql.py

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -32,31 +32,31 @@ class SchemaNotFoundError(FileNotFoundError):
3232
class Connection(object):
3333
def __init__(self, autoconnect=True, verbose=False, **kwargs):
3434
"""
35-
Set up the connection to the WRDS database.
36-
By default, also establish the connection to the database.
37-
38-
Optionally, the user may specify connection parameters:
39-
*wrds_hostname*: WRDS database hostname
40-
*wrds_port*: database connection port number
41-
*wrds_dbname*: WRDS database name
42-
*wrds_username*: WRDS username
43-
*autoconnect*: If false will not immediately establish the connection
44-
45-
The constructor will use the .pgpass file if it exists and may make use of
46-
PostgreSQL environment variables such as PGHOST, PGUSER, etc., if cooresponding
47-
parameters are not set.
48-
If not, it will ask the user for a username and password.
49-
It will also direct the user to information on setting up .pgpass.
50-
51-
Additionally, creating the instance will load a list of schemas
52-
the user has permission to access.
53-
54-
:return: None
55-
56-
Usage::
57-
>>> db = wrds.Connection()
58-
Loading library list...
59-
Done
35+
Set up the connection to the WRDS database.
36+
By default, also establish the connection to the database.
37+
38+
Optionally, the user may specify connection parameters:
39+
*wrds_hostname*: WRDS database hostname
40+
*wrds_port*: database connection port number
41+
*wrds_dbname*: WRDS database name
42+
*wrds_username*: WRDS username
43+
*autoconnect*: If false will not immediately establish the connection
44+
45+
The constructor will use the .pgpass file if it exists and may make use of
46+
PostgreSQL environment variables such as PGHOST, PGUSER, etc., if cooresponding
47+
parameters are not set.
48+
If not, it will ask the user for a username and password.
49+
It will also direct the user to information on setting up .pgpass.
50+
51+
Additionally, creating the instance will load a list of schemas
52+
the user has permission to access.
53+
54+
:return: None
55+
56+
Usage::
57+
>>> db = wrds.Connection()
58+
Loading library list...
59+
Done
6060
"""
6161
self._verbose = verbose
6262
self._password = ""
@@ -74,7 +74,7 @@ def __init__(self, autoconnect=True, verbose=False, **kwargs):
7474
self.connect()
7575
self.load_library_list()
7676

77-
def __make_sa_engine_conn(self, raise_err = False):
77+
def __make_sa_engine_conn(self, raise_err=False):
7878
username = self._username
7979
hostname = self._hostname
8080
password = urllib.parse.quote_plus(self._password)
@@ -127,7 +127,7 @@ def connect(self):
127127
try:
128128
self.create_pgpass_file()
129129
print("Created .pgpass file successfully.")
130-
except:
130+
except Exception:
131131
print("Failed to create .pgpass file.")
132132
print(
133133
"You can create this file yourself at any time "
@@ -216,7 +216,7 @@ def __get_user_credentials(self):
216216

217217
def create_pgpass_file(self):
218218
"""
219-
Create a .pgpass file to store WRDS connection credentials..
219+
Create a .pgpass file to store WRDS connection credentials.
220220
221221
Use the existing username and password if already connected to WRDS,
222222
or prompt for that information if not.

‎wrds/test.py

Lines changed: 79 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import wrds
44
import unittest
5+
56
try:
67
import unittest.mock as mock
78
except ImportError:
@@ -10,134 +11,147 @@
1011

1112

1213
class TestInitMethod(unittest.TestCase):
13-
""" Test the wrds.Connection.__init__() method,
14-
with both default and custom parameters.
14+
"""Test the wrds.Connection.__init__() method,
15+
with both default and custom parameters.
1516
"""
16-
@mock.patch('wrds.sql.sa')
17+
18+
@mock.patch("wrds.sql.sa")
1719
def test_init_calls_sqlalchemy_create_engine_defaults(self, mock_sa):
1820
wrds.Connection()
19-
connstring = 'postgresql://{host}:{port}/{dbname}'
21+
connstring = "postgresql://{host}:{port}/{dbname}"
2022
connstring = connstring.format(
2123
host=wrds.sql.WRDS_POSTGRES_HOST,
2224
port=wrds.sql.WRDS_POSTGRES_PORT,
23-
dbname=wrds.sql.WRDS_POSTGRES_DB)
25+
dbname=wrds.sql.WRDS_POSTGRES_DB,
26+
)
2427
mock_sa.create_engine.assert_called_with(
2528
connstring,
26-
connect_args={'sslmode': 'require',
27-
'application_name': wrds.sql.appname},
28-
isolation_level='AUTOCOMMIT')
29+
connect_args={"sslmode": "require", "application_name": wrds.sql.appname},
30+
isolation_level="AUTOCOMMIT",
31+
)
2932

30-
@mock.patch('wrds.sql.sa')
33+
@mock.patch("wrds.sql.sa")
3134
def test_init_calls_sqlalchemy_create_engine_custom(self, mock_sa):
32-
username = 'faketestusername'
33-
connstring = 'postgresql://{usr}@{host}:{port}/{dbname}'
35+
username = "faketestusername"
36+
connstring = "postgresql://{usr}@{host}:{port}/{dbname}"
3437
connstring = connstring.format(
3538
usr=username,
3639
host=wrds.sql.WRDS_POSTGRES_HOST,
3740
port=wrds.sql.WRDS_POSTGRES_PORT,
38-
dbname=wrds.sql.WRDS_POSTGRES_DB)
41+
dbname=wrds.sql.WRDS_POSTGRES_DB,
42+
)
3943
wrds.Connection(wrds_username=username)
4044
mock_sa.create_engine.assert_called_with(
4145
connstring,
42-
connect_args={'sslmode': 'require',
43-
'application_name': wrds.sql.appname},
44-
isolation_level='AUTOCOMMIT')
46+
connect_args={"sslmode": "require", "application_name": wrds.sql.appname},
47+
isolation_level="AUTOCOMMIT",
48+
)
4549

46-
@mock.patch('wrds.sql.Connection.load_library_list')
47-
@mock.patch('wrds.sql.Connection.connect')
50+
@mock.patch("wrds.sql.Connection.load_library_list")
51+
@mock.patch("wrds.sql.Connection.connect")
4852
def test_init_default_connect(self, mock_connect, mock_lll):
4953
wrds.Connection()
5054
mock_connect.assert_called_once()
5155

52-
@mock.patch('wrds.sql.Connection.connect')
56+
@mock.patch("wrds.sql.Connection.connect")
5357
def test_init_autoconnect_false_no_connect(self, mock_connect):
5458
wrds.Connection(autoconnect=False)
5559
mock_connect.assert_not_called()
5660

57-
@mock.patch('wrds.sql.Connection.connect')
58-
@mock.patch('wrds.sql.Connection.load_library_list')
61+
@mock.patch("wrds.sql.Connection.connect")
62+
@mock.patch("wrds.sql.Connection.load_library_list")
5963
def test_init_default_load_library_list(self, mock_lll, mock_connect):
6064
wrds.Connection()
6165
mock_lll.assert_called_once()
6266

63-
@mock.patch('wrds.sql.Connection.connect')
64-
@mock.patch('wrds.sql.Connection.load_library_list')
65-
def test_init_autoconnect_false_no_connect_second_function(self, mock_lll, mock_connect):
67+
@mock.patch("wrds.sql.Connection.connect")
68+
@mock.patch("wrds.sql.Connection.load_library_list")
69+
def test_init_autoconnect_false_no_connect_second_function(
70+
self, mock_lll, mock_connect
71+
):
6672
wrds.Connection(autoconnect=False)
6773
mock_lll.assert_not_called()
6874

6975

7076
class TestConnectMethod(unittest.TestCase):
71-
""" Test the wrds.Connection.connect method.
77+
"""
78+
Test the wrds.Connection.connect method.
7279
73-
Since all exceptions are caught immediately,
74-
I'm just not smart enough to simulate bad passwords with
75-
the code as written.
80+
Since all exceptions are caught immediately,
81+
I'm just not smart enough to simulate bad passwords with
82+
the code as written.
7683
"""
7784

7885
def setUp(self):
7986
self.t = wrds.Connection(autoconnect=False)
80-
self.t._hostname = 'wrds.test.private'
87+
self.t._hostname = "wrds.test.private"
8188
self.t._port = 12345
82-
self.t._username = 'faketestusername'
83-
self.t._password = 'faketestuserpass'
84-
self.t._dbname = 'testdbname'
89+
self.t._username = "faketestusername"
90+
self.t._password = "faketestuserpass"
91+
self.t._dbname = "testdbname"
8592
self.t._Connection__get_user_credentials = mock.Mock()
86-
self.t._Connection__get_user_credentials.return_value = (self.t._username, self.t._password)
93+
self.t._Connection__get_user_credentials.return_value = (
94+
self.t._username,
95+
self.t._password,
96+
)
8797

8898
def test_connect_calls_sqlalchemy_engine_connect(self):
8999
self.t.engine = mock.Mock()
90100
self.t.connect()
91101
self.t.engine.connect.assert_called_once()
92102

93-
@mock.patch('wrds.sql.sa')
103+
@mock.patch("wrds.sql.sa")
94104
def test_connect_calls_get_user_credentials_on_exception(self, mock_sa):
95105
self.t.engine = mock.Mock()
96-
self.t.engine.connect.side_effect = Exception('Fake exception for testing')
106+
self.t.engine.connect.side_effect = Exception("Fake exception for testing")
97107
self.t.connect()
98108
self.t._Connection__get_user_credentials.assert_called_once()
99109

100-
@mock.patch('wrds.sql.sa')
110+
@mock.patch("wrds.sql.sa")
101111
def test_connect_calls_sqlalchemy_create_engine_on_exception(self, mock_sa):
102112
self.t.engine = mock.Mock()
103-
self.t.engine.connect.side_effect = Exception('Fake exception for testing')
104-
connstring = 'postgresql://{usr}:{pwd}@{host}:{port}/{dbname}'
113+
self.t.engine.connect.side_effect = Exception("Fake exception for testing")
114+
connstring = "postgresql://{usr}:{pwd}@{host}:{port}/{dbname}"
105115
connstring = connstring.format(
106116
usr=self.t._username,
107117
pwd=self.t._password,
108118
host=self.t._hostname,
109119
port=self.t._port,
110-
dbname=self.t._dbname)
120+
dbname=self.t._dbname,
121+
)
111122
self.t.connect()
112123
mock_sa.create_engine.assert_called_with(
113124
connstring,
114-
connect_args={'sslmode': 'require',
115-
'application_name': wrds.sql.appname},
116-
isolation_level='AUTOCOMMIT')
125+
connect_args={"sslmode": "require", "application_name": wrds.sql.appname},
126+
isolation_level="AUTOCOMMIT",
127+
)
117128

118129

119130
class TestRawSqlMethod(unittest.TestCase):
120-
""" Test the wrds.Connection.raw_sql method.
131+
"""Test the wrds.Connection.raw_sql method.
121132
122-
wrds.Connection.raw_sql() should be able to take
123-
'normal' and parameterized SQL,
124-
and throw an error if not all parameters are supplied.
133+
wrds.Connection.raw_sql() should be able to take
134+
'normal' and parameterized SQL,
135+
and throw an error if not all parameters are supplied.
125136
"""
126137

127138
def setUp(self):
128139
self.t = wrds.Connection(autoconnect=False)
129-
self.t._hostname = 'wrds.test.private'
140+
self.t._hostname = "wrds.test.private"
130141
self.t._port = 12345
131-
self.t._username = 'faketestusername'
132-
self.t._password = 'faketestuserpass'
133-
self.t._dbname = 'testdbname'
142+
self.t._username = "faketestusername"
143+
self.t._password = "faketestuserpass"
144+
self.t._dbname = "testdbname"
134145
self.t._Connection__get_user_credentials = mock.Mock()
135-
self.t._Connection__get_user_credentials.return_value = (self.t._username, self.t._password)
146+
self.t._Connection__get_user_credentials.return_value = (
147+
self.t._username,
148+
self.t._password,
149+
)
136150
self.t.connection = mock.Mock()
137151
self.t.engine = mock.Mock()
138152

139-
@mock.patch('wrds.sql.sa')
140-
@mock.patch('wrds.sql.pd')
153+
@mock.patch("wrds.sql.sa")
154+
@mock.patch("wrds.sql.pd")
141155
def test_rawsql_takes_unparameterized_sql(self, mock_pd, mock_sa):
142156
sql = "SELECT * FROM information_schema.tables LIMIT 1"
143157
self.t.raw_sql(sql)
@@ -150,10 +164,13 @@ def test_rawsql_takes_unparameterized_sql(self, mock_pd, mock_sa):
150164
params=None,
151165
)
152166

153-
@mock.patch('wrds.sql.sa')
154-
@mock.patch('wrds.sql.pd')
167+
@mock.patch("wrds.sql.sa")
168+
@mock.patch("wrds.sql.pd")
155169
def test_rawsql_takes_parameterized_sql(self, mock_pd, mock_sa):
156-
sql = "SELECT * FROM information_schema.tables where table_name = %(tablename)s LIMIT 1"
170+
sql = (
171+
"SELECT * FROM information_schema.tables "
172+
"WHERE table_name = %(tablename)s LIMIT 1"
173+
)
157174
tablename = "pg_stat_activity"
158175
self.t.engine = mock.Mock()
159176
self.t.raw_sql(sql, params=tablename)
@@ -171,7 +188,10 @@ class TestCreatePgpassFile(unittest.TestCase):
171188
def setUp(self):
172189
self.t = wrds.Connection(autoconnect=False)
173190
self.t._Connection__get_user_credentials = mock.Mock()
174-
self.t._Connection__get_user_credentials.return_value = ('faketestusername', 'faketestpassword')
191+
self.t._Connection__get_user_credentials.return_value = (
192+
"faketestusername",
193+
"faketestpassword",
194+
)
175195
self.t._Connection__create_pgpass_file_win32 = mock.Mock()
176196
self.t._Connection__create_pgpass_file_unix = mock.Mock()
177197

@@ -185,16 +205,16 @@ def test_create_pgpass_calls_get_user_credentials_if_not_password(self):
185205
self.t.create_pgpass_file()
186206
self.t._Connection__get_user_credentials.assert_called_once()
187207

188-
@unittest.skipIf(sys.platform != 'win32', 'Windows-only test')
208+
@unittest.skipIf(sys.platform != "win32", "Windows-only test")
189209
def test_create_pgpass_calls_win32_version_if_windows(self):
190210
self.t.create_pgpass_file()
191211
self.t._Connection__create_pgpass_file_win32.assert_called_once()
192212

193-
@unittest.skipIf(sys.platform == 'win32', 'Unix-only test')
213+
@unittest.skipIf(sys.platform == "win32", "Unix-only test")
194214
def test_create_pgpass_calls_unix_version_if_unix(self):
195215
self.t.create_pgpass_file()
196216
self.t._Connection__create_pgpass_file_unix.assert_called_once()
197217

198218

199-
if (__name__ == '__main__'):
219+
if __name__ == "__main__":
200220
unittest.main()

0 commit comments

Comments
 (0)