@@ -4639,19 +4639,23 @@ def select(self, *args: Hashable | list[Hashable]):
4639
4639
1 Cooper 22
4640
4640
2 Marley 35
4641
4641
"""
4642
- if args and isinstance (args [0 ], list ):
4642
+ err_msg = (
4643
+ "`DataFrame.select` supports individual columns "
4644
+ "`df.select('col1', 'col2',...)` or a list "
4645
+ "`df.select(['col1', 'col2',...])`, but not both. "
4646
+ "You can unpack the list if you have a mix: "
4647
+ "`df.select(*['col1', 'col2'], 'col3')`."
4648
+ )
4649
+ list_or_star_args = list (args )
4650
+ if args and isinstance (list_or_star_args [0 ], list ):
4643
4651
if len (args ) == 1 :
4644
- columns = args [0 ]
4652
+ columns = list_or_star_args [0 ]
4645
4653
else :
4646
- raise ValueError (
4647
- "`DataFrame.select` supports individual columns "
4648
- "`df.select('col1', 'col2',...)` or a list "
4649
- "`df.select(['col1', 'col2',...])`, but not both. "
4650
- "You can unpack the list if you have a mix: "
4651
- "`df.select(*['col1', 'col2'], 'col3')`."
4652
- )
4654
+ raise ValueError (err_msg )
4655
+ elif any (isinstance (arg , list ) for arg in args ):
4656
+ raise ValueError (err_msg )
4653
4657
else :
4654
- columns = list ( args )
4658
+ columns = list_or_star_args # type: ignore[assignment]
4655
4659
4656
4660
indexer = self .columns ._get_indexer_strict (columns , "columns" )[1 ]
4657
4661
return self .take (indexer , axis = 1 )
0 commit comments