Update to have a consistent style

This commit is contained in:
Ryan Leckey
2014-12-12 10:05:34 -08:00
parent 3a2991b221
commit 8e36f3d6b1
2 changed files with 62 additions and 31 deletions

View File

@@ -236,7 +236,8 @@ class EncryptedType(TypeDecorator, ScalarCoercible):
try: try:
value = self.underlying_type.process_bind_param( value = self.underlying_type.process_bind_param(
value, dialect) value, dialect
)
except AttributeError: except AttributeError:
# Doesn't have 'process_bind_param' # Doesn't have 'process_bind_param'
@@ -244,7 +245,7 @@ class EncryptedType(TypeDecorator, ScalarCoercible):
# Handle 'boolean' and 'dates' # Handle 'boolean' and 'dates'
type_ = self.underlying_type.python_type type_ = self.underlying_type.python_type
if issubclass(type_, bool): if issubclass(type_, bool):
value = "true" if value else "false" value = 'true' if value else 'false'
elif issubclass(type_, (datetime.date, datetime.time)): elif issubclass(type_, (datetime.date, datetime.time)):
value = value.isoformat() value = value.isoformat()
@@ -259,7 +260,8 @@ class EncryptedType(TypeDecorator, ScalarCoercible):
try: try:
return self.underlying_type.process_result_value( return self.underlying_type.process_result_value(
decrypted_value, dialect) decrypted_value, dialect
)
except AttributeError: except AttributeError:
# Doesn't have 'process_result_value' # Doesn't have 'process_result_value'
@@ -267,19 +269,22 @@ class EncryptedType(TypeDecorator, ScalarCoercible):
# Handle 'boolean' and 'dates' # Handle 'boolean' and 'dates'
type_ = self.underlying_type.python_type type_ = self.underlying_type.python_type
if issubclass(type_, bool): if issubclass(type_, bool):
return decrypted_value == "true" return decrypted_value == 'true'
elif issubclass(type_, datetime.datetime): elif issubclass(type_, datetime.datetime):
return datetime.datetime.strptime( return datetime.datetime.strptime(
decrypted_value, "%Y-%m-%dT%H:%M:%S") decrypted_value, '%Y-%m-%dT%H:%M:%S'
)
elif issubclass(type_, datetime.time): elif issubclass(type_, datetime.time):
return datetime.datetime.strptime( return datetime.datetime.strptime(
decrypted_value, "%H:%M:%S").time() decrypted_value, '%H:%M:%S'
).time()
elif issubclass(type_, datetime.date): elif issubclass(type_, datetime.date):
return datetime.datetime.strptime( return datetime.datetime.strptime(
decrypted_value, "%Y-%m-%d").date() decrypted_value, '%Y-%m-%d'
).date()
# Handle all others # Handle all others
return self.underlying_type.python_type(decrypted_value) return self.underlying_type.python_type(decrypted_value)

View File

@@ -16,7 +16,7 @@ from sqlalchemy_utils.types.encrypted import AesEngine, FernetEngine
@mark.skipif('cryptography is None') @mark.skipif('cryptography is None')
class EncryptedTypeTestCase(TestCase): class EncryptedTypeTestCase(TestCase):
@pytest.fixture(scope="function") @pytest.fixture(scope='function')
def user(self, request): def user(self, request):
# set the values to the user object # set the values to the user object
self.user = self.User() self.user = self.User()
@@ -49,7 +49,7 @@ class EncryptedTypeTestCase(TestCase):
def generate_test_token(self): def generate_test_token(self):
import string import string
import random import random
token = "" token = ''
characters = string.ascii_letters + string.digits characters = string.ascii_letters + string.digits
for i in range(60): for i in range(60):
token += ''.join(random.choice(characters)) token += ''.join(random.choice(characters))
@@ -61,7 +61,7 @@ class EncryptedTypeTestCase(TestCase):
self.user_name = u'someone' self.user_name = u'someone'
self.user_phone = u'(555) 555-5555' self.user_phone = u'(555) 555-5555'
self.user_color = u'#fff' self.user_color = u'#fff'
self.user_enum = "One" self.user_enum = 'One'
self.user_date = date(2010, 10, 2) self.user_date = date(2010, 10, 2)
self.user_time = time(10, 12) self.user_time = time(10, 12)
self.user_datetime = datetime(2010, 10, 2, 10, 12) self.user_datetime = datetime(2010, 10, 2, 10, 12)
@@ -73,46 +73,66 @@ class EncryptedTypeTestCase(TestCase):
class User(self.Base): class User(self.Base):
__tablename__ = 'user' __tablename__ = 'user'
id = sa.Column(sa.Integer, primary_key=True) id = sa.Column(sa.Integer, primary_key=True)
username = sa.Column(EncryptedType( username = sa.Column(EncryptedType(
sa.Unicode, sa.Unicode,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
access_token = sa.Column(EncryptedType( access_token = sa.Column(EncryptedType(
sa.String, sa.String,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
is_active = sa.Column(EncryptedType( is_active = sa.Column(EncryptedType(
sa.Boolean, sa.Boolean,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
accounts_num = sa.Column(EncryptedType( accounts_num = sa.Column(EncryptedType(
sa.Integer, sa.Integer,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
phone = sa.Column(EncryptedType( phone = sa.Column(EncryptedType(
PhoneNumberType, PhoneNumberType,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
color = sa.Column(EncryptedType( color = sa.Column(EncryptedType(
ColorType, ColorType,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
date = sa.Column(EncryptedType( date = sa.Column(EncryptedType(
sa.Date, sa.Date,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
time = sa.Column(EncryptedType( time = sa.Column(EncryptedType(
sa.Time, sa.Time,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
datetime = sa.Column(EncryptedType( datetime = sa.Column(EncryptedType(
sa.DateTime, sa.DateTime,
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
enum = sa.Column(EncryptedType( enum = sa.Column(EncryptedType(
sa.Enum("One", name="user_enum_t"), sa.Enum('One', name='user_enum_t'),
self.test_key, self.test_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
self.User = User self.User = User
@@ -123,7 +143,8 @@ class EncryptedTypeTestCase(TestCase):
name = sa.Column(EncryptedType( name = sa.Column(EncryptedType(
sa.Unicode, sa.Unicode,
lambda: self._team_key, lambda: self._team_key,
self.__class__.encryption_engine)) self.__class__.encryption_engine)
)
self.Team = Team self.Team = Team
@@ -159,14 +180,14 @@ class EncryptedTypeTestCase(TestCase):
def test_lookup_key(self): def test_lookup_key(self):
# Add teams # Add teams
self._team_key = "one" self._team_key = 'one'
team = self.Team(key=self._team_key, name="One") team = self.Team(key=self._team_key, name='One')
self.session.add(team) self.session.add(team)
self.session.flush() self.session.flush()
team_1_id = team.id team_1_id = team.id
self._team_key = "two" self._team_key = 'two'
team = self.Team(key=self._team_key, name="Two") team = self.Team(key=self._team_key, name='Two')
self.session.add(team) self.session.add(team)
self.session.flush() self.session.flush()
team_2_id = team.id team_2_id = team.id
@@ -175,19 +196,23 @@ class EncryptedTypeTestCase(TestCase):
# Lookup teams # Lookup teams
self._team_key = self.session.query(self.Team.key).filter_by( self._team_key = self.session.query(self.Team.key).filter_by(
id=team_1_id).one()[0] id=team_1_id
).one()[0]
team = self.session.query(self.Team).get(team_1_id) team = self.session.query(self.Team).get(team_1_id)
assert team.name == "One" assert team.name == 'One'
with pytest.raises(Exception): with pytest.raises(Exception):
self.session.query(self.Team).get(team_2_id) self.session.query(self.Team).get(team_2_id)
self._team_key = self.session.query(self.Team.key).filter_by( self._team_key = self.session.query(self.Team.key).filter_by(
id=team_2_id).one()[0] id=team_2_id
).one()[0]
team = self.session.query(self.Team).get(team_2_id) team = self.session.query(self.Team).get(team_2_id)
assert team.name == "Two" assert team.name == 'Two'
with pytest.raises(Exception): with pytest.raises(Exception):
self.session.query(self.Team).get(team_1_id) self.session.query(self.Team).get(team_1_id)
@@ -203,7 +228,8 @@ class TestAesEncryptedTypeTestcase(EncryptedTypeTestCase):
def test_lookup_by_encrypted_string(self, user): def test_lookup_by_encrypted_string(self, user):
test = self.session.query(self.User).filter( test = self.session.query(self.User).filter(
self.User.username == self.user_name).first() self.User.username == self.user_name
).first()
assert test.username == user.username assert test.username == user.username