datamodel-code-generator icon indicating copy to clipboard operation
datamodel-code-generator copied to clipboard

Generating collection methods to RootModel

Open alkshmir opened this issue 1 year ago • 10 comments
trafficstars

Is your feature request related to a problem? Please describe. When I convert an array like following with datamodel-codegen --input openapi.yaml --output model.py --output-model-type pydantic_v2.BaseModel,

openapi: 3.0.3

components:
  schemas:
    MyItem:
      type: object
      properties:
        foo:
          type: number
        bar:
          type: string

    MyArray:
      type: array
      items:
        $ref: "#/components/schemas/MyItem"

MyArray cannot be used like a list because it lacks methods that collections should have.

For example, the following code prints ('root', [MyItem(foo=123.0, bar='text')]) instead of MyItem(foo=123.0, bar='text').

from model import MyArray, MyItem

arr = MyArray([MyItem(foo=123, bar="text")])

for item in arr:
    print(item)

Describe the solution you'd like It would be nice to generate the methods that collections should have, e.g., __getitem__, __iter__, or __len__:

class MyItem(BaseModel):
    foo: Optional[float] = None
    bar: Optional[str] = None


class MyArray(RootModel[List[MyItem]]):
    root: List[MyItem]

    def __getitem__(self, index):
        return self.root[index]

    def __iter__(self):
        return iter(self.root)

    def __len__(self):
         return len(self.root)

This would be a common use case so it would be great that some CLI option enables this feature.

Describe alternatives you've considered Writing a collection methods manually on every generated RootModel, or pydantic might have some other pretty solutions.

Additional context Other collection methods like __setitem__, __delitem__, __contains__, or some other methods list-like collections have (append, extend, insert, remove, pop, clear, or index) might be better to generated as well.

Sorry if this request was a duplicate. Thank you very much for the great tool.

alkshmir avatar Feb 01 '24 08:02 alkshmir

for item in arr.root:
    print(item)

rafalkrupinski avatar Feb 01 '24 20:02 rafalkrupinski

Thanks for your help. Accessing root directly can be a solution but I'm just wondering how to use the model in the same way as List[MyItem] not to depend on the internal structure.

alkshmir avatar Feb 01 '24 22:02 alkshmir

Technically root is a part of pydantics public API, like model_extra and other properties, but I understand it may be an inconvenience.

rafalkrupinski avatar Feb 02 '24 06:02 rafalkrupinski

I did not check the discussion but this may help me (but it seems they are talking about pydantic v1 and idk it will work for v2 as well): https://github.com/koxudaxi/datamodel-code-generator/discussions/640

alkshmir avatar Feb 06 '24 00:02 alkshmir

@rafalkrupinski Thank you for your help.

@alkshmir

As your link shows, this would be a good solution to use a custom template and run the CLI to generate your own root model. What do you think?

koxudaxi avatar Mar 16 '24 16:03 koxudaxi