ソースを参照

Allow load & dump from yaml

theenglishway (time) 6 年 前
コミット
4c7eb92e87
5 ファイル変更60 行追加2 行削除
  1. 2 1
      .gitignore
  2. 16 0
      pyplanner/cli.py
  3. 30 0
      pyplanner/database.py
  4. 11 0
      pyplanner/models/__init__.py
  5. 1 1
      setup.cfg

+ 2 - 1
.gitignore

@@ -2,4 +2,5 @@
 *.egg-info
 venv/
 *.pyc
-*.db
+*.db
+*.yaml

+ 16 - 0
pyplanner/cli.py

@@ -14,6 +14,22 @@ def main(ctx, db_url):
     ctx.obj['session'] = db.session_factory()
 
 
+@main.command()
+@click.option('-f', '--file', type=str, default='./data.yaml', show_default=True)
+@click.pass_context
+def dump(ctx, file):
+    db = ctx.obj['db']
+    db.dump(file)
+
+
+@main.command()
+@click.option('-f', '--file', type=str, default='./data.yaml', show_default=True)
+@click.pass_context
+def load(ctx, file):
+    db = ctx.obj['db']
+    db.load(file)
+
+
 @main.group()
 @click.pass_context
 def milestone(ctx):

+ 30 - 0
pyplanner/database.py

@@ -1,15 +1,45 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.ext.declarative import declarative_base
+import inspect
+from ruamel import yaml
 
 
 SQLABase = declarative_base()
 
 
+yaml = yaml.YAML()
+yaml.default_flow_style = False
+yaml.allow_unicode = True
+
 class Database:
     def __init__(self, db_url):
         self.engine = create_engine(db_url)
         self.session_factory = sessionmaker(bind=self.engine)
 
         SQLABase.metadata.create_all(self.engine)
+        self.models = {
+            k: v for k, v in SQLABase._decl_class_registry.items()
+            if inspect.isclass(v) and issubclass(v, SQLABase)
+        }
+
+    def dump(self, file):
+        session = self.session_factory()
+        instances_dict = {}
+
+        for k, v in self.models.items():
+            instances = [i for i in session.query(v).all()]
+            instances_dict[k] = instances
+            yaml.register_class(v)
+
+        with open(file, 'w') as f:
+            yaml.dump(instances_dict, f)
+
+    def load(self, file):
+        with open(file) as f:
+            instances_dict = yaml.load(f)
 
+        session = self.session_factory()
+        for k, v_list in instances_dict.items():
+            model = self.models[k]
+            instances = [model(**v) for v in v_list]

+ 11 - 0
pyplanner/models/__init__.py

@@ -17,6 +17,17 @@ class Base:
         except ValidationError as e:
             return False, e
 
+    @classmethod
+    def to_yaml(cls, representer, node):
+        print(cls, node)
+        return representer.represent_data(
+            {k: v for k, v in node.__dict__.items() if not k.startswith('_')}
+        )
+
+    @classmethod
+    def from_yaml(cls, constructor, node):
+        return cls()
+
 
 from .milestones import Milestone
 from .sprints import Sprint

+ 1 - 1
setup.cfg

@@ -25,7 +25,7 @@ install_requires =
 	Cerberus
 	Click
 	colorama
-	PyYaml
+	ruamel.yaml
 	termcolor
 	TinyDb