diff --git a/pydbapi/col/colmodel.py b/pydbapi/col/colmodel.py index 3a9dada..1d64dfd 100644 --- a/pydbapi/col/colmodel.py +++ b/pydbapi/col/colmodel.py @@ -2,9 +2,11 @@ # @Author: longfengpili # @Date: 2023-06-02 15:27:41 # @Last Modified by: longfengpili -# @Last Modified time: 2023-07-27 15:33:37 +# @Last Modified time: 2023-10-31 10:25:06 # @github: https://github.com/longfengpili +from typing import Iterable, List, Any + class ColumnModel(object): @@ -49,34 +51,47 @@ def create_sqlexpr(self): class ColumnsModel(object): def __init__(self, *columns): - self.columns = columns + self.columns = list(columns) def __repr__(self): return f"{self.columns}" - def __getitem__(self, key): - return self.columns[key] + def __getitem__(self, index: int): + if isinstance(index, slice): + start, stop, step = index.indices(len(self.columns)) + return [self.columns[i] for i in range(start, stop, step)] + elif 0 <= index < len(self.columns): + return self.columns[index] + else: + raise IndexError("Index out of range") + + def __setitem__(self, index: int, column: ColumnModel): + if 0 <= index < len(self.columns): + self.columns[index] = column + else: + raise IndexError("Index out of range") + + def append(self, item: Any) -> None: + self.columns.append(item) + + def extend(self, iterable: Iterable[Any]) -> None: + self.columns.extend(iterable) + + def __len__(self) -> int: + return len(self.columns) # def __iter__(self): # for column in self.columns: # yield column - def index(self, key): - return self.columns[key] - def __contains__(self, name): col = self.get_column_by_name(name) isin = True if col else False return isin - def append(self, column: ColumnModel): - columns = list(self.columns) - columns.append(column) - return ColumnsModel(*columns) - def remove(self, remove_column: str): new_columns = [] - columns = list(self.columns) + columns = self.columns for column in columns: if column.newname != remove_column: new_columns.append(column) diff --git a/tests/colmodel/test_colmodel.py b/tests/colmodel/test_colmodel.py index 2312960..e7d8a8c 100644 --- a/tests/colmodel/test_colmodel.py +++ b/tests/colmodel/test_colmodel.py @@ -2,7 +2,7 @@ # @Author: longfengpili # @Date: 2023-06-02 15:27:41 # @Last Modified by: longfengpili -# @Last Modified time: 2023-07-27 15:41:50 +# @Last Modified time: 2023-10-31 10:26:10 # @github: https://github.com/longfengpili @@ -24,12 +24,24 @@ def teardown_method(self, method): def test_column_getter(self): colname = 'score' - col = self.columns.index(0) - print(col) col = self.columns.get_column_by_name(colname) print(col) col = self.columns[0] print(col) + cols = self.columns[1:3] + print(cols) + + def test_column_append(self): + test = ColumnModel('test', 'varchar', desc='test') + self.columns.append(test) + print(self.columns) + + def test_column_extend(self): + test = ColumnModel('test', 'varchar', desc='test') + test2 = ColumnModel('test2', 'varchar', desc='test2') + tests = ColumnsModel(test, test2) + self.columns.extend(tests) + print(self.columns) def test_column_contain(self): colname = 'score'