@@ -25,6 +25,10 @@ type LeetcodeResponse struct {
2525			Tags        []struct  {
2626				Name  string  `json:"name"` 
2727			} `json:"topicTags"` 
28+ 			CodeSnippets  []struct  {
29+ 				Lang  string  `json:"lang"` 
30+ 				Code  string  `json:"code"` 
31+ 			} `json:"codeSnippets"` 
2832		} `json:"question"` 
2933	} `json:"data"` 
3034}
@@ -117,6 +121,7 @@ type ReadmeData struct {
117121	Difficulty  string 
118122	Tags        []string 
119123	Problem     string 
124+ 	Code        string 
120125	Now         string 
121126}
122127
@@ -157,6 +162,10 @@ func fetchLeetcode(slug string) (*ReadmeData, error) {
157162    topicTags { 
158163      name 
159164    } 
165+     codeSnippets { 
166+       lang 
167+       code 
168+     } 
160169  } 
161170}` 
162171
@@ -184,13 +193,22 @@ func fetchLeetcode(slug string) (*ReadmeData, error) {
184193		tags  =  append (tags , tag .Name )
185194	}
186195
196+ 	code  :=  "" 
197+ 	for  _ , snippet  :=  range  q .CodeSnippets  {
198+ 		if  snippet .Lang  ==  "Go"  ||  snippet .Lang  ==  "Golang"  {
199+ 			code  =  snippet .Code 
200+ 			break 
201+ 		}
202+ 	}
203+ 
187204	return  & ReadmeData {
188205		ID :         q .QuestionID ,
189206		Title :      strings .ReplaceAll (q .Title , " " , "-" ),
190207		Slug :       slug ,
191208		Difficulty : q .Difficulty ,
192209		Tags :       tags ,
193210		Problem :    stripHTML (q .Content ),
211+ 		Code :       code ,
194212		Now :        time .Now ().Format (time .RFC3339 ),
195213	}, nil 
196214}
@@ -228,6 +246,38 @@ func writeREADME(data *ReadmeData) error {
228246	return  tmpl .Execute (f , data )
229247}
230248
249+ func  writeMainGo (data  * ReadmeData ) error  {
250+ 	dir  :=  fmt .Sprintf ("Leetcode/%s.%s" , data .ID , data .Title )
251+ 	f , err  :=  os .Create (filepath .Join (dir , "main.go" ))
252+ 	if  err  !=  nil  {
253+ 		return  err 
254+ 	}
255+ 	defer  f .Close ()
256+ 	_ , err  =  f .WriteString (data .Code )
257+ 	return  err 
258+ }
259+ 
260+ func  writeMainTest (data  * ReadmeData ) error  {
261+ 	dir  :=  fmt .Sprintf ("Leetcode/%s.%s" , data .ID , data .Title )
262+ 	f , err  :=  os .Create (filepath .Join (dir , "main_test.go" ))
263+ 	if  err  !=  nil  {
264+ 		return  err 
265+ 	}
266+ 	defer  f .Close ()
267+ 	testContent  :=  `package main 
268+ 
269+ import ( 
270+ 	"testing" 
271+ ) 
272+ 
273+ func TestExample(t *testing.T) { 
274+ 	// TODO: add test cases 
275+ } 
276+ ` 
277+ 	_ , err  =  f .WriteString (testContent )
278+ 	return  err 
279+ }
280+ 
231281func  main () {
232282	if  len (os .Args ) <  2  {
233283		fmt .Println ("Usage: leetcode-readme-gen <id>" )
@@ -252,5 +302,11 @@ func main() {
252302	if  err  :=  writeREADME (data ); err  !=  nil  {
253303		panic (err )
254304	}
255- 	fmt .Printf ("README generated for %s (%s)\n " , data .Title , data .ID )
305+ 	if  err  :=  writeMainGo (data ); err  !=  nil  {
306+ 		panic (err )
307+ 	}
308+ 	if  err  :=  writeMainTest (data ); err  !=  nil  {
309+ 		panic (err )
310+ 	}
311+ 	fmt .Printf ("README, main.go, and main_test.go generated for %s (%s)\n " , data .Title , data .ID )
256312}
0 commit comments