Support CSV (#308)

* Support CSV * ignore types * fix name * remove unused code * fix invalid condition on input-file-type
koxudaxi · Jan 18, 2021 · dffdc4f · dffdc4f
1 parent eb4c28e
commit dffdc4f
Show file tree

Hide file tree

Showing 7 changed files with 104 additions and 16 deletions.
diff --git a/README.md b/README.md
@@ -17,7 +17,7 @@ See [documentation](https://koxudaxi.github.io/datamodel-code-generator) for mor
 ## Supported source types
 - OpenAPI 3 (YAML/JSON, [OpenAPI Data Type](https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.2.md#data-types))
 - JSON Schema ([JSON Schema Core](http://json-schema.org/draft/2019-09/json-schema-validation.html)/[JSON Schema Validation](http://json-schema.org/draft/2019-09/json-schema-validation.html))
-- JSON/YAML Data (it will be converted to JSON Schema)
+- JSON/YAML/CSV Data (it will be converted to JSON Schema)
 - Python dictionary (it will be converted to JSON Schema)
 
 ## Implemented list
@@ -67,7 +67,7 @@ $ pip install datamodel-code-generator[http]
 The `datamodel-codegen` command:
 ```
 usage: datamodel-codegen [-h] [--input INPUT]
-                         [--input-file-type {auto,openapi,jsonschema,json,yaml,dict}]
+                         [--input-file-type {auto,openapi,jsonschema,json,yaml,dict,csv}]
                          [--output OUTPUT] [--base-class BASE_CLASS] [--field-constraints]
                          [--snake-case-field] [--strip-default-none]
                          [--allow-population-by-field-name] [--use-default] [--force-optional]
@@ -82,7 +82,7 @@ usage: datamodel-codegen [-h] [--input INPUT]
 optional arguments:
   -h, --help            show this help message and exit
   --input INPUT         Input file/directory (default: stdin)
-  --input-file-type {auto,openapi,jsonschema,json,yaml,dict}
+  --input-file-type {auto,openapi,jsonschema,json,yaml,dict,csv}
                         Input file type (default: auto)
   --output OUTPUT       Output file (default: stdout)
   --base-class BASE_CLASS

diff --git a/datamodel_code_generator/__init__.py b/datamodel_code_generator/__init__.py
@@ -68,8 +68,6 @@ def __get__(self, instance: Any, owner: Any = None) -> Any:
 DEFAULT_BASE_CLASS: str = 'pydantic.BaseModel'
 
 
-# ALL_MODEL: str = '#all#'
-
 SafeLoader.yaml_constructors[
     'tag:yaml.org,2002:timestamp'
 ] = SafeLoader.yaml_constructors['tag:yaml.org,2002:str']
@@ -140,6 +138,7 @@ class InputFileType(Enum):
     Json = 'json'
     Yaml = 'yaml'
     Dict = 'dict'
+    CSV = 'csv'
 
 
 class Error(Exception):
@@ -219,20 +218,36 @@ def generate(
 
         parser_class = JsonSchemaParser
 
-        if input_file_type in [
+        if input_file_type in (
             InputFileType.Json,
             InputFileType.Yaml,
             InputFileType.Dict,
-        ]:
+            InputFileType.CSV,
+        ):
             try:
                 if isinstance(input_, Path) and input_.is_dir():  # pragma: no cover
                     raise Error(f'Input must be a file for {input_file_type}')
-                input_text = (
-                    input_
-                    if isinstance(input_, str)
-                    else input_.read_text(encoding=encoding)
-                )
-                obj: Dict[Any, Any] = load_yaml(input_text)
+                obj: Dict[Any, Any]
+                if input_file_type == InputFileType.CSV:
+                    import csv
+
+                    def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
+                        csv_reader = csv.DictReader(csv_file)
+                        return dict(zip(csv_reader.fieldnames, next(csv_reader)))  # type: ignore
+
+                    if isinstance(input_, str):
+                        import io
+
+                        obj = get_header_and_first_line(io.StringIO(input_))
+                    else:
+                        with input_.open(encoding=encoding) as f:
+                            obj = get_header_and_first_line(f)
+                else:
+                    obj = load_yaml(
+                        input_
+                        if isinstance(input_, str)
+                        else input_.read_text(encoding=encoding)
+                    )
             except:
                 raise Error('Invalid file format')
             from genson import SchemaBuilder

diff --git a/docs/index.md b/docs/index.md
@@ -13,7 +13,7 @@ This code generator creates pydantic model from an openapi file and others.
 ## Supported source types
 - OpenAPI 3 (YAML/JSON, [OpenAPI Data Type](https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.2.md#data-types))
 - JSON Schema ([JSON Schema Core](http://json-schema.org/draft/2019-09/json-schema-validation.html) /[JSON Schema Validation](http://json-schema.org/draft/2019-09/json-schema-validation.html))
-- JSON/YAML Data (it will be converted to JSON Schema)
+- JSON/YAML/CSV Data (it will be converted to JSON Schema)
 - Python dictionary (it will be converted to JSON Schema)
 
 ## Installation
@@ -34,7 +34,7 @@ $ pip install datamodel-code-generator[http]
 The `datamodel-codegen` command:
 ```
 usage: datamodel-codegen [-h] [--input INPUT]
-                         [--input-file-type {auto,openapi,jsonschema,json,yaml,dict}]
+                         [--input-file-type {auto,openapi,jsonschema,json,yaml,dict,csv}]
                          [--output OUTPUT] [--base-class BASE_CLASS] [--field-constraints]
                          [--snake-case-field] [--strip-default-none]
                          [--allow-population-by-field-name] [--use-default] [--force-optional]
@@ -49,7 +49,7 @@ usage: datamodel-codegen [-h] [--input INPUT]
 optional arguments:
   -h, --help            show this help message and exit
   --input INPUT         Input file/directory (default: stdin)
-  --input-file-type {auto,openapi,jsonschema,json,yaml,dict}
+  --input-file-type {auto,openapi,jsonschema,json,yaml,dict,csv}
                         Input file type (default: auto)
   --output OUTPUT       Output file (default: stdout)
   --base-class BASE_CLASS

diff --git a/tests/data/csv/simple.csv b/tests/data/csv/simple.csv
@@ -0,0 +1,4 @@
+id,name,tel,zip code
+1,taro,0123456789,98765
+2,ken,234567891,98764
+3,ichiro,345678912,98763
diff --git a/tests/data/expected/main/csv_file_simple/output.py b/tests/data/expected/main/csv_file_simple/output.py
@@ -0,0 +1,14 @@
+# generated by datamodel-codegen:
+#   filename:  simple.csv
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class Model(BaseModel):
+    id: str
+    name: str
+    tel: str
+    zip_code: str = Field(..., alias='zip code')
diff --git a/tests/data/expected/main/csv_stdin_simple/output.py b/tests/data/expected/main/csv_stdin_simple/output.py
@@ -0,0 +1,14 @@
+# generated by datamodel-codegen:
+#   filename:  <stdin>
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class Model(BaseModel):
+    id: str
+    name: str
+    tel: str
+    zip_code: str = Field(..., alias='zip code')
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -18,6 +18,7 @@
 JSON_DATA_PATH: Path = DATA_PATH / 'json'
 YAML_DATA_PATH: Path = DATA_PATH / 'yaml'
 PYTHON_DATA_PATH: Path = DATA_PATH / 'python'
+CSV_DATA_PATH: Path = DATA_PATH / 'csv'
 EXPECTED_MAIN_PATH = DATA_PATH / 'expected' / 'main'
 
 TIMESTAMP = '1985-10-26T01:21:00-07:00'
@@ -1371,3 +1372,43 @@ def test_space_and_special_characters_dict():
         )
     with pytest.raises(SystemExit):
         main()
+
+
+@freeze_time('2019-07-26')
+def test_csv_file():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(CSV_DATA_PATH / 'simple.csv'),
+                '--output',
+                str(output_file),
+                '--input-file-type',
+                'csv',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (EXPECTED_MAIN_PATH / 'csv_file_simple' / 'output.py').read_text()
+        )
+    with pytest.raises(SystemExit):
+        main()
+
+
+@freeze_time('2019-07-26')
+def test_csv_stdin(monkeypatch):
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        monkeypatch.setattr('sys.stdin', (CSV_DATA_PATH / 'simple.csv').open())
+        return_code: Exit = main(
+            ['--output', str(output_file), '--input-file-type', 'csv',]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (EXPECTED_MAIN_PATH / 'csv_stdin_simple' / 'output.py').read_text()
+        )
+    with pytest.raises(SystemExit):
+        main()