@@ -143,12 +143,19 @@ def _is_file_or_dir(self, cloud_path: S3Path) -> Optional[str]:
143143 return "dir"
144144
145145 def _exists (self , cloud_path : S3Path ) -> bool :
146+ # check if this is a bucket
147+ if not cloud_path .key :
148+ try :
149+ self .client .head_bucket (Bucket = cloud_path .bucket )
150+ return True
151+ except ClientError :
152+ return False
153+
146154 return self ._s3_file_query (cloud_path ) is not None
147155
148156 def _s3_file_query (self , cloud_path : S3Path ):
149157 """Boto3 query used for quick checks of existence and if path is file/dir"""
150- # first check if this is an object that we can access directly
151-
158+ # check if this is an object that we can access directly
152159 try :
153160 obj = self .s3 .Object (cloud_path .bucket , cloud_path .key )
154161 obj .load ()
@@ -169,40 +176,55 @@ def _s3_file_query(self, cloud_path: S3Path):
169176 )
170177
171178 def _list_dir (self , cloud_path : S3Path , recursive = False ) -> Iterable [Tuple [S3Path , bool ]]:
172- bucket = self .s3 .Bucket (cloud_path .bucket )
173-
174179 prefix = cloud_path .key
175180 if prefix and not prefix .endswith ("/" ):
176181 prefix += "/"
177182
178183 yielded_dirs = set ()
179184
180- if recursive :
181- for o in bucket .objects .filter (Prefix = prefix ):
182- # get directory from this path
183- for parent in PurePosixPath (o .key [len (prefix ) :]).parents :
184- # if we haven't surfaced their directory already
185- if parent not in yielded_dirs and str (parent ) != "." :
186- yield (self .CloudPath (f"s3://{ cloud_path .bucket } { prefix } { parent } ), True )
187- yielded_dirs .add (parent )
185+ paginator = self .client .get_paginator ("list_objects_v2" )
188186
189- yield (self .CloudPath (f"s3://{ o .bucket_name } { o .key } ), False )
190- else :
191- # non recursive is best done with old client API rather than resource
192- paginator = self .client .get_paginator ("list_objects" )
193-
194- for result in paginator .paginate (
195- Bucket = cloud_path .bucket , Prefix = prefix , Delimiter = "/"
196- ):
197- # sub directory names
198- for result_prefix in result .get ("CommonPrefixes" , []):
187+ for result in paginator .paginate (
188+ Bucket = cloud_path .bucket , Prefix = prefix , Delimiter = ("" if recursive else "/" )
189+ ):
190+ # yield everything in common prefixes as directories
191+ for result_prefix in result .get ("CommonPrefixes" , []):
192+ canonical = result_prefix .get ("Prefix" ).rstrip ("/" ) # keep a canonical form
193+ if canonical not in yielded_dirs :
194+ yield (
195+ self .CloudPath (f"s3://{ cloud_path .bucket } { canonical } ),
196+ True ,
197+ )
198+ yielded_dirs .add (canonical )
199+
200+ # check all the keys
201+ for result_key in result .get ("Contents" , []):
202+ # yield all the parents of any key that have not been yielded already
203+ o_relative_path = result_key .get ("Key" )[len (prefix ) :]
204+ for parent in PurePosixPath (o_relative_path ).parents :
205+ parent_canonical = prefix + str (parent ).rstrip ("/" )
206+ if parent_canonical not in yielded_dirs and str (parent ) != "." :
207+ yield (
208+ self .CloudPath (f"s3://{ cloud_path .bucket } { parent_canonical } ),
209+ True ,
210+ )
211+ yielded_dirs .add (parent_canonical )
212+
213+ # if we already yielded this dir, go to next item in contents
214+ canonical = result_key .get ("Key" ).rstrip ("/" )
215+ if canonical in yielded_dirs :
216+ continue
217+
218+ # s3 fake directories have 0 size and end with "/"
219+ if result_key .get ("Key" ).endswith ("/" ) and result_key .get ("Size" ) == 0 :
199220 yield (
200- self .CloudPath (f"s3://{ cloud_path .bucket } { result_prefix . get ( 'Prefix' ) } ),
221+ self .CloudPath (f"s3://{ cloud_path .bucket } { canonical } ),
201222 True ,
202223 )
224+ yielded_dirs .add (canonical )
203225
204- # files in the directory
205- for result_key in result . get ( "Contents" , []) :
226+ # yield object as file
227+ else :
206228 yield (
207229 self .CloudPath (f"s3://{ cloud_path .bucket } { result_key .get ('Key' )} ),
208230 False ,
0 commit comments