Merge pull request #4 from crashappsec/enums

adding --import/--only flags

.gitignore

@@ -1,2 +1,5 @@
 Pipfile*
 *__pycache__*
+tags
+build
+*.egg-info

gql_schema_codegen/__init__.py

@@ -1,6 +1,6 @@
 from .schema.schema import Schema
 
 VERSION = (1, 0, 0)
-__version__ = '.'.join(str(x) for x in VERSION)
+__version__ = ".".join(str(x) for x in VERSION)
 
-__all__ = ['Schema']
+__all__ = ["Schema"]

gql_schema_codegen/__main__.py

@@ -3,6 +3,7 @@
 
 from .schema import Schema
 
+
 parser = argparse.ArgumentParser(
     description="Generate python file with types from a  GraphQL schema file."
 )
@@ -24,7 +25,6 @@
     "-t",
     dest="to_path",
     type=str,
-    default="generated_schema_types.py",
     help="wanted output file path (default: schema_types.py)",
 )
 
@@ -37,19 +37,52 @@
     help="path of the config file in yaml format (default: gql_schema_codegen.config.yml)",
 )
 
+blocks = {
+    "enum",
+    "scalar",
+    "type",
+    "param_type",
+    "input",
+}
+
+parser.add_argument(
+    "-b",
+    "--block",
+    dest="blocks",
+    action="append",
+    choices=blocks,
+    help="operate on these blocks differently",
+)
+
+group = parser.add_mutually_exclusive_group()
+
+group.add_argument(
+    "--import",
+    dest="import_blocks",
+    help="module name from where to import --blocks",
+)
+
+group.add_argument(
+    "--only",
+    dest="only_blocks",
+    action="store_true",
+    help="only output blocks specified in --blocks",
+)
+
 
 args = parser.parse_args()
 
 if __name__ == "__main__":
-    schema_path, schema_url, to_path, config_file = (
-        args.schema_path,
-        args.schema_url,
-        args.to_path,
-        args.config_file,
+    s = Schema(
+        path=args.schema_path,
+        url=args.schema_url,
+        config_file=args.config_file,
+        blocks=args.blocks,
+        import_blocks=args.import_blocks,
+        only_blocks=args.only_blocks,
     )
-
-    s = Schema(path=schema_path, url=schema_url, config_file=config_file)
-    file_representation = s.generate_type_file(to_path)
-
-    sys.stdout.write(file_representation)
-    sys.stdout.close()
+    if args.to_path:
+        file_representation = s.generate_type_file(args.to_path)
+    else:
+        sys.stdout.write(s.file_representation)
+        sys.stdout.close()

gql_schema_codegen/base/__init__.py

@@ -1,3 +1,3 @@
 from .base import BaseInfo
 
-__all__ = ['BaseInfo']
+__all__ = ["BaseInfo"]

gql_schema_codegen/base/base.py

@@ -2,7 +2,6 @@
 
 
 class BaseInfo:
-
     def __init__(self, info: NamedTuple) -> None:
         self.info = info
 
@@ -11,4 +10,5 @@ def __getattr__(self, item):
             return getattr(self.info, item)
 
         raise AttributeError(
-            f"{self.__class__.__name__} object has no attribute {item}")
+            f"{self.__class__.__name__} object has no attribute {item}"
+        )

gql_schema_codegen/block/__init__.py

@@ -1,3 +1,3 @@
 from .block import Block, BlockInfo, BlockField, BlockFieldInfo
 
-__all__ = ['Block', 'BlockInfo', 'BlockField', 'BlockFieldInfo']
+__all__ = ["Block", "BlockInfo", "BlockField", "BlockFieldInfo"]

gql_schema_codegen/block/block.py

@@ -25,12 +25,18 @@ def __init__(self, info, dependency_group: DependencyGroup) -> None:
         self.parent_classes = set()
 
     @property
-    def heading_file_line(self):
+    def display_name(self):
         display_name = pascal_case(self.name)
 
         if self.type == "param_type":
             display_name = f"{display_name}Params"
 
+        return display_name
+
+    @property
+    def heading_file_line(self):
+        display_name = self.display_name
+
         if self.type == "enum":
             self.dependency_group.add_dependency(
                 Dependency(imported_from="enum", dependency="Enum")

gql_schema_codegen/constants/__init__.py

@@ -1,4 +1,21 @@
-from .constants import DIRECTIVE_PATTERN, DIRECTIVE_USAGE_PATTERN, SCALAR_PATTERN, UNION_PATTERN, VALUE_TYPES, BLOCK_PATTERN, FIELD_PATTERN, RESOLVER_TYPES
+from .constants import (
+    DIRECTIVE_PATTERN,
+    DIRECTIVE_USAGE_PATTERN,
+    SCALAR_PATTERN,
+    UNION_PATTERN,
+    VALUE_TYPES,
+    BLOCK_PATTERN,
+    FIELD_PATTERN,
+    RESOLVER_TYPES,
+)
 
-__all__ = ['DIRECTIVE_PATTERN', 'DIRECTIVE_USAGE_PATTERN', 'SCALAR_PATTERN',
-           'UNION_PATTERN', 'VALUE_TYPES', 'BLOCK_PATTERN', 'FIELD_PATTERN', 'RESOLVER_TYPES']
+__all__ = [
+    "DIRECTIVE_PATTERN",
+    "DIRECTIVE_USAGE_PATTERN",
+    "SCALAR_PATTERN",
+    "UNION_PATTERN",
+    "VALUE_TYPES",
+    "BLOCK_PATTERN",
+    "FIELD_PATTERN",
+    "RESOLVER_TYPES",
+]

gql_schema_codegen/dependency/__init__.py

@@ -1,3 +1,3 @@
 from .dependency import Dependency, DependencyGroup
 
-__all__ = ['Dependency', 'DependencyGroup']
+__all__ = ["Dependency", "DependencyGroup"]

gql_schema_codegen/scalar/__init__.py

@@ -1,3 +1,3 @@
 from .scalar import ScalarType, ScalarInfo
 
-__all__ = ['ScalarType', 'ScalarInfo']
+__all__ = ["ScalarType", "ScalarInfo"]

gql_schema_codegen/scalar/scalar.py

@@ -29,7 +29,7 @@ def file_representation(self):
             )
             self.dependency_group.add_direct_dependency("dateutil.parser")
             self.dependency_group.add_direct_dependency("neo4j")
-            return ""
+            return f"{self.name} = datetime"
 
         if type(self.value) is not str:
             self.dependency_group.add_dependency(

gql_schema_codegen/schema/__init__.py

@@ -1,3 +1,3 @@
 from .schema import Schema
 
-__all__ = ['Schema']
+__all__ = ["Schema"]

gql_schema_codegen/schema/schema.py

@@ -2,17 +2,27 @@
 import os
 import re
 import subprocess
-from typing import List, Optional
+from typing import List, Optional, Set
 
 import yaml
-from graphql import (build_client_schema, build_schema,
-                     get_introspection_query, print_schema)
+from graphql import (
+    build_client_schema,
+    build_schema,
+    get_introspection_query,
+    print_schema,
+)
 from graphqlclient import GraphQLClient
 
 from ..block import Block, BlockField, BlockFieldInfo, BlockInfo
-from ..constants import (BLOCK_PATTERN, DIRECTIVE_PATTERN,
-                         DIRECTIVE_USAGE_PATTERN, FIELD_PATTERN,
-                         RESOLVER_TYPES, SCALAR_PATTERN, UNION_PATTERN)
+from ..constants import (
+    BLOCK_PATTERN,
+    DIRECTIVE_PATTERN,
+    DIRECTIVE_USAGE_PATTERN,
+    FIELD_PATTERN,
+    RESOLVER_TYPES,
+    SCALAR_PATTERN,
+    UNION_PATTERN,
+)
 from ..constants.block_fields import all_block_fields
 from ..dependency import Dependency, DependencyGroup
 from ..scalar import ScalarInfo, ScalarType
@@ -26,12 +36,15 @@ class Schema:
     _cleaned_string: Optional[str] = None
     _unions: List[UnionType] = []
     _scalars: List[ScalarType] = []
-    ignore_types = []
+    ignore_types: List[str] = []
     generate_empty_params_types = False
     path: Optional[str] = None
     url: Optional[str] = None
     _config_file_content = None
     config_file: Optional[str] = None
+    _special_blocks: Set[str] = set()
+    _import_blocks: Optional[str] = None
+    _only_blocks: bool = False
 
     def __init__(self, **kwargs) -> None:
         if "path" in kwargs and type(kwargs["path"]) is str:
@@ -43,6 +56,10 @@ def __init__(self, **kwargs) -> None:
         if "config_file" in kwargs and type(kwargs["config_file"]) is str:
             self.config_file = kwargs["config_file"]
 
+        self._special_blocks = kwargs.get("blocks", self._special_blocks)
+        self._import_blocks = kwargs.get("import_blocks", self._import_blocks)
+        self._only_blocks = kwargs.get("only_blocks", self._only_blocks)
+
         self.dependency_group = DependencyGroup()
 
     @property
@@ -192,6 +209,7 @@ def blocks(self):
 
                 block_type = block["type"]
                 block_name = block["name"]
+
                 all_block_fields[block_name] = set()
                 for field in self.get_fields_from_block(block["fields"]):
                     all_block_fields[block_name].add(field["name"])
@@ -215,7 +233,14 @@ def blocks(self):
                 )
                 b = Block(block_info, dependency_group=self.dependency_group)
 
-                blocks.append(b)
+                if self._import_blocks and block_type in self._special_blocks:
+                    self.dependency_group.add_dependency(
+                        Dependency(self._import_blocks, b.display_name)
+                    )
+                elif self._only_blocks and block_type not in self._special_blocks:
+                    continue
+                else:
+                    blocks.append(b)
 
             self._blocks = list(
                 filter(lambda b: b.name not in self.ignore_types, blocks)
@@ -253,9 +278,18 @@ def scalars(self):
                 scalar_info = ScalarInfo(
                     name=u["name"], value=self.custom_scalars.get(u["name"])
                 )
-                self._scalars.append(
-                    ScalarType(scalar_info, dependency_group=self.dependency_group)
-                )
+                scalar = ScalarType(scalar_info, dependency_group=self.dependency_group)
+                if self._import_blocks and "scalar" in self._special_blocks:
+                    # hack to adjust dependency group
+                    scalar.file_representation
+                    self.dependency_group.add_dependency(
+                        Dependency(self._import_blocks, scalar_info.name)
+                    )
+                else:
+                    self._scalars.append(scalar)
+
+        if self._only_blocks and "scalar" not in self._special_blocks:
+            return []
 
         return self._scalars
 

gql_schema_codegen/union/__init__.py

@@ -1,3 +1,3 @@
 from .union import UnionType, UnionInfo
 
-__all__ = ['UnionType', 'UnionInfo']
+__all__ = ["UnionType", "UnionInfo"]

gql_schema_codegen/union/union.py

@@ -10,18 +10,17 @@ class UnionInfo(NamedTuple):
 
 
 class UnionType(BaseInfo):
-
     def __init__(self, info: UnionInfo, dependency_group: DependencyGroup) -> None:
         super().__init__(info)
         self.dependency_group = dependency_group
 
     @property
     def separated_types(self):
-        return re.sub(r'[^\w\|]', '', self.types).split('|')
+        return re.sub(r"[^\w\|]", "", self.types).split("|")
 
     @property
     def cleaned_types(self):
-        return ', '.join(map(lambda t: f"'{t}'", self.separated_types))
+        return ", ".join(map(lambda t: f"'{t}'", self.separated_types))
 
     @property
     def file_representation(self):

gql_schema_codegen/utils/__init__.py

@@ -1,3 +1,3 @@
 from .pascal_case import pascal_case
 
-__all__ = ['pascal_case']
+__all__ = ["pascal_case"]

gql_schema_codegen/utils/pascal_case.py

@@ -3,8 +3,11 @@
 
 def pascal_case(string: str):
     cleaned_string = string.replace(" ", "")
-    first_letter = re.search(r'\w', cleaned_string)
+    first_letter = re.search(r"\w", cleaned_string)
     if not first_letter:
-        return ''
+        return ""
 
-    return cleaned_string[:first_letter.start() + 1].upper() + cleaned_string[first_letter.start() + 1:]
+    return (
+        cleaned_string[: first_letter.start() + 1].upper()
+        + cleaned_string[first_letter.start() + 1 :]
+    )

setup.cfg

@@ -0,0 +1,2 @@
+[flake8]
+ignore = E501

setup.py

@@ -1,4 +1,5 @@
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
+
 
 packages = find_packages()
 
@@ -15,14 +16,18 @@
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/sauldom102/gql_schema_codegen",
-    py_modules=['gql_schema_codegen', ],
+    py_modules=[
+        "gql_schema_codegen",
+    ],
     install_requires=[
-        'graphql-core==3.2.0',
-        'graphqlclient==0.2.4',
-        'argparse==1.4.0',
-        'pytest-snapshot==0.8.1',
-        'pytest==7.1.1',
-        'PyYAML==6.0'
+        "graphql-core>=3.2.0",
+        "graphqlclient>=0.2.4",
+        "neo4j",
+        "pyyaml",
+    ],
+    test_require=[
+        "pytest",
+        "pytest-snapshot",
     ],
     classifiers=[
         "Programming Language :: Python :: 3",