7dracoder
diff --git a/‎spanner_orm/admin/migration_executor.py‎
Lines changed: 24 additions & 16 deletions b/‎spanner_orm/admin/migration_executor.py‎
Lines changed: 24 additions & 16 deletions
diff --git a/‎spanner_orm/admin/migration_manager.py‎
Lines changed: 12 additions & 7 deletions b/‎spanner_orm/admin/migration_manager.py‎
Lines changed: 12 additions & 7 deletions
diff --git a/‎spanner_orm/admin/scripts.py‎
Lines changed: 7 additions & 2 deletions b/‎spanner_orm/admin/scripts.py‎
Lines changed: 7 additions & 2 deletions
@@ -13,8 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Handles execution of migrations."""
+
+from __future__ import annotations
+
 import datetime
 import logging
+from typing import Any, Dict, Optional
 
 from spanner_orm import api
 from spanner_orm import error
@@ -24,34 +28,37 @@
 from spanner_orm.admin import migration_status
 from spanner_orm.admin import update
 
+from google.auth import credentials as auth_credentials
+
 _logger = logging.getLogger(__name__)
 
 
-class MigrationExecutor(object):
+class MigrationExecutor:
   """Handles execution of migrations."""
 
   def __init__(self,
-               instance,
-               database,
-               project=None,
-               credentials=None,
-               basedir=None):
+               instance: str,
+               database: str,
+               project: Optional[str] = None,
+               credentials: Optional[auth_credentials.Credentials] = None,
+               basedir: Optional[str] = None):
     self._manager = migration_manager.MigrationManager(basedir)
     self._migration_status_map = None
     self._instance = instance
     self._database = database
     self._project = project
     self._credentials = credentials
 
-  def migrated(self, migration_id):
+  def migrated(self, migration_id: str) -> bool:
     if migration_id is None:
       return True
     return self._migration_status().get(migration_id, False)
 
-  def migrations(self):
+  # TODO(dbrandao): make a Migration object so this is no longer Any
+  def migrations(self) -> Any:
     return self._manager.migrations
 
-  def migrate(self, target_migration=None):
+  def migrate(self, target_migration: Optional[str] = None) -> None:
     """Executes unmigrated migrations on the curent database.
 
     Note: SpannerApi and SpannerAdminApi connections are modified as a result
@@ -79,7 +86,7 @@ def migrate(self, target_migration=None):
       self._update_status(migration.migration_id, True)
     self._hangup()
 
-  def rollback(self, target_migration):
+  def rollback(self, target_migration: str) -> None:
     """Rolls back migrated migrations on the curent database.
 
     Note: SpannerApi and SpannerAdminApi connections are modified as a result
@@ -110,7 +117,7 @@ def rollback(self, target_migration):
       self._update_status(migration.migration_id, False)
     self._hangup()
 
-  def _connect(self):
+  def _connect(self) -> None:
     admin_api.SpannerAdminApi.connect(
         self._instance,
         self._database,
@@ -122,11 +129,12 @@ def _connect(self):
         project=self._project,
         credentials=self._credentials)
 
-  def _hangup(self):
+  def _hangup(self) -> None:
     admin_api.SpannerAdminApi.hangup()
     api.SpannerApi.hangup()
 
-  def _filter_migrations(self, migrations, migrated, last_migration):
+  def _filter_migrations(self, migrations: Any, migrated: bool,
+                         last_migration: Optional[str]) -> Any:
     """Filters the list of migrations according to the desired conditions.
 
     Args:
@@ -153,7 +161,7 @@ def _filter_migrations(self, migrations, migrated, last_migration):
               last_migration))
     return filtered
 
-  def _migration_status(self):
+  def _migration_status(self) -> Dict[str, bool]:
     """Gathers from Spanner which migrations have been executed."""
     if self._migration_status_map is None:
       model_from_db = metadata.SpannerMetadata.model(
@@ -167,7 +175,7 @@ def _migration_status(self):
 
     return self._migration_status_map
 
-  def _update_status(self, migration_id, new_status):
+  def _update_status(self, migration_id: str, new_status: bool) -> None:
     """Updates migration status in the database for the given migration."""
     new_model = migration_status.MigrationStatus({
         'id': migration_id,
@@ -178,7 +186,7 @@ def _update_status(self, migration_id, new_status):
         None, [new_model], force_write=True)
     self._migration_status()[migration_id] = new_status
 
-  def _validate_migrations(self):
+  def _validate_migrations(self) -> None:
     """Validates the migration status of all migrations makes sense."""
     migrations = self.migrations()
     if not migrations:
 
@@ -13,28 +13,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Handles reading and writing of migration files."""
+
+from __future__ import annotations
+
+
 import datetime
 import importlib
 import os
 import re
 import string
+from typing import Any, Iterable, List, Optional
 import uuid
 
 from spanner_orm import error
 
 
-class MigrationManager(object):
+class MigrationManager:
   """Handles reading and writing of migration files."""
   DEFAULT_DIRECTORY = 'migrations'
 
-  def __init__(self, basedir=None):
+  def __init__(self, basedir: Optional[str] = None):
     self.basedir = basedir or self.DEFAULT_DIRECTORY
     self._migrations = None
 
     if not os.path.exists(self.basedir):
       os.makedirs(self.basedir)
 
-  def generate(self, migration_name):
+  def generate(self, migration_name: str) -> str:
     """Creates a new migration that is the last migration to be executed."""
     migration_id = uuid.uuid4().hex[-12:]
     prev_id = self.migrations[-1].migration_id if self.migrations else None
@@ -58,14 +63,14 @@ def generate(self, migration_name):
     return filepath
 
   @property
-  def migrations(self):
+  def migrations(self) -> Any:
     """Loads and orders all migrations in the base dir."""
     if self._migrations is None:
       unordered_migrations = self._all_migrations()
       self._migrations = self._order_migrations(unordered_migrations)
     return self._migrations
 
-  def _migration_from_file(self, filename):
+  def _migration_from_file(self, filename: str) -> Any:
     """Loads a single migration from the given filename in the base dir."""
     module_name = re.sub(r'\W', '_', filename)
     path = os.path.join(self.basedir, filename)
@@ -76,7 +81,7 @@ def _migration_from_file(self, filename):
       raise error.SpannerError('{} has no migration id'.format(path))
     return module
 
-  def _all_migrations(self):
+  def _all_migrations(self) -> List[Any]:
     """Loads all migrations from the base dir."""
     migrations = []
     for filename in os.listdir(self.basedir):
@@ -85,7 +90,7 @@ def _all_migrations(self):
         migrations.append(self._migration_from_file(filename))
     return migrations
 
-  def _order_migrations(self, migrations):
+  def _order_migrations(self, migrations: Iterable[Any]) -> List[Any]:
     """Returns list of migrations in the order they have to be applied."""
     if not migrations:
       return []
 
@@ -13,16 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Entry point for spanner_orm scripts."""
+
+from __future__ import annotations
+
 import argparse
+from typing import Any
+
 from spanner_orm.admin import migration_manager
 
 
-def generate(args):
+def generate(args: Any) -> None:
   manager = migration_manager.MigrationManager(args.directory)
   manager.generate(args.name)
 
 
-def main(as_module=False):
+def main(as_module: bool = False) -> None:
   prog = 'spanner-orm' if as_module else None
   parser = argparse.ArgumentParser(prog=prog)
   subparsers = parser.add_subparsers(