From my blog: https://david.coffee/refactoring-code-with-chat-gpt/
ChatGPT is really damn impressive and the more I read about it, the more it feels like withcraft. I see people using it for game development, coming up with titles, improving prompts for MidJourney and and and. It's shaping up to be one of those universally useful tools that are great to keep around.
Here's one thing I had on my mind: Can it help me when I code?
I was working on a little pet project and had some pretty inefficient code I slapped together for a middleware while watching TV:
package middleware
import (
"context"
"encoding/json"
"net/http"
)
var parsedJsonBodyContextKey struct{}
func ParseJsonBody(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var data map[string]interface{}
err := json.NewDecoder(r.Body).Decode(&data)
if err != nil {
ctx := context.WithValue(r.Context(), parsedJsonBodyContextKey, nil)
next.ServeHTTP(w, r.WithContext(ctx))
return
}
ctx := context.WithValue(r.Context(), parsedJsonBodyContextKey, nil)
next.ServeHTTP(w, r.WithContext(ctx))
})
}
The code takes whatever is passed into the request body, tries to JSON decode it, then puts the result into context. If no body exists or isn't JSON, it will pass nil
instead.
I wanted to shorten and make it a bit more efficient but though it'd be a cool experiment to let ChatGPT do it:
Me: Can you refactor the following code to make it more efficient?
<code from above here>
ChatGPT: Here are a few ways the provided Go code can be imporved:
- The
ParseJsonBody()
function defines adata
variable to store the parsed JSON body, but it doesn't appear to be used anywhere in the function. This variable can be removed.- The
ParseJsonBody()
function creates a newcontext.Context
with a key-value pair for the parsed JSON body, but it sets the value tonil
in both branches of theif
statement. This means that the value of the parsed JSON body will always benil
, regardless of whether the JSON body was parsed successfully or not. The value of the parsed JSON body should be set to the parsed JSON data if parsing was successful, andnil
if parsing failed.- The
ParseJsonBody()
function passes the originalhttp.Request
object to thenext
handler, even if the JSON body was parsed successfully. In this case, thehttp.Request
object should be updated with the newcontext.Context
that contains the parsed JSON body before being passed to thenext
handler.Here is how the code can be improved:
package middleware import ( "context" "encoding/json" "net/http" ) var parsedJsonBodyContextKey struct{} func ParseJsonBody(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var data map[string]interface{} err := json.NewDecoder(r.Body).Decode(&data) if err != nil { ctx := context.WithValue(r.Context(), parsedJsonBodyContextKey, nil) next.ServeHTTP(w, r.WithContext(ctx)) return } ctx := context.WithValue(r.Context(), parsedJsonBodyContextKey, data) next.ServeHTTP(w, r.WithContext(ctx)) }) }
Huh, turns out I copy pasted and forgot to change the normal branch to actually add data
into the context - whoops. ChatGPT spotted this, understood what this was supposed to do and fixed it.
Can we make it more compact though?
Me: Can you make the code more compact and more DRY?
ChatGPT: Yes, the provided code can be made more compact and DRY (Don't Repeat Yourself) by removing the duplicate code and using a single
context.Context
object.Here is how the code can be improved:
package middleware import ( "context" "encoding/json" "net/http" ) var parsedJsonBodyContextKey struct{} func ParseJsonBody(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var data map[string]interface{} err := json.NewDecoder(r.Body).Decode(&data) ctx := context.WithValue(r.Context(), parsedJsonBodyContextKey, data) next.ServeHTTP(w, r.WithContext(ctx)) }) }
This is exactly what I had in mind for refactoring this code: No need for the if condition at all because var data map[string]interface{}
is nil
by default. Only nitpick would be that we assign err
and never use it, so the compiler will complain about this code.
Let's add a function to get the data out of the context again. Can ChatGPT summon this function for us?
Me: Please add a function that takes the body out of the context again
ChatGPT: To extract the parsed JSON body from the
context.Context
object, a new function can be added to themiddleware
package. This function can take acontext.Context
object as an argument and return the parsed JSON body, if it exists in the context, ornil
otherwise.Here is an example implementation of this function:
package middleware import ( "context" "encoding/json" ) var parsedJsonBodyContextKey struct{} func ParseJsonBody(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var data map[string]interface{} err := json.NewDecoder(r.Body).Decode(&data) ctx := context.WithValue(r.Context(), parsedJsonBodyContextKey, data) next.ServeHTTP(w, r.WithContext(ctx)) }) } func GetParsedJsonBody(ctx context.Context) map[string]interface{} { val := ctx.Value(parsedJsonBodyContextKey) if val == nil { return nil } return val.(map[string]interface{}) }
To use this function, you can pass a
context.Context
object as an argument and check the return value fornil
to determine whether the JSON body was parsed successfully or not.Here is an example of how this function can be used:
package main import ( "context" "fmt" "net/http" "github.com/myproject/middleware" ) func main() { http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { parsedBody := middleware.GetParsedJsonBody(r.Context()) if parsedBody == nil { w.Write([]byte("JSON body was not parsed")) return } w.Write([]byte(fmt.Sprintf("
Short and concise. It understood what I wanted it to do, it delivered a solution and even gave an explanation with examples how to use it!
After a bit of hacking I noticed that using map[string]interface{}
is kind of a PITA because we have to do many type assertions and checks if keys exist. I got curious if we can tell it to swap out the type into something custom that has a nicer interface:
Me: Instead of using map[string]interface{}, let's use a custom type that has methods for getting the value for a given key, or a default
ChatGPT: To use a custom type that has methods for getting the value for a given key or a default value, you can define a new
ParsedJsonBody
type that satisfies thecontext.Context
interface. This type can be a struct that contains amap[string]interface{}
field to store the parsed JSON body.Here is an example implementation of the
ParsedJsonBody
type:type ParsedJsonBody struct { data map[string]interface{} } func (p *ParsedJsonBody) Value(key interface{}) interface{} { if key == parsedJsonBodyContextKey { return p } return nil } func (p *ParsedJsonBody) Get(key string) interface{} { return p.data[key] } func (p *ParsedJsonBody) GetString(key string) string { val := p.data[key] if val == nil { return "" } return val.(string) } func (p *ParsedJsonBody) GetInt(key string) int { val := p.data[key] if val == nil { return 0 } return val.(int) }
The
ParseJsonBody()
middleware function can be updated to use theParsedJsonBody
type instead ofmap[string]interface{}
:
Hmm okay, this is where things get a bit weird, but that could have also been my way of explaining what I wanted.
The code looks ok on first skimming, but it ignored my instruction to have a default. I think it interpreted that as "we'll return a default", but not one I can specify.
There's lots of type assertions in there that don't have any checks, so this will panic if we were to try to GetInt
something that isn't an int. Also, you can define a new ParsedJsonBody
type that satisfies the context.Context
interface is nonsense, and so is the Value()
method.
What I actually wanted was a custom type like type MyType map[string]interface{}
, not a struct with a type field, but whatever.
Let's ask directly if it thinks anything is wrong with that code:
Me: Thanks for the code, do you see any problems with this implementation?
ChatGPT: There is a potential problem with the provided implementation of the
ParsedJsonBody
type. TheValue()
method returns a pointer to theParsedJsonBody
struct, but theGet()
,GetString()
, andGetInt()
methods return values of the underlyingdata
field. This means that if theParsedJsonBody
struct is modified after it is added to thecontext.Context
object, the modifications will be reflected in thecontext.Context
, but not in the return values of theGet()
,GetString()
, andGetInt()
methods.To fix this problem, the
Get()
,GetString()
, andGetInt()
methods should return copies of the values in thedata
field, rather than the values themselves.Here is an updated implementation of the
ParsedJsonBody
type that fixes this problem:type ParsedJsonBody struct { data map[string]interface{} } func (p *ParsedJsonBody) Value(key interface{}) interface{} { if key == parsedJsonBodyContextKey { return p } return nil } func (p *ParsedJsonBody) Get(key string) interface{} { val, ok := p.data[key] if !ok { return nil } return val } func (p *ParsedJsonBody) GetString(key string) string { val, ok := p.data[key].(string) if !ok { return "" } return val } func (p *ParsedJsonBody) GetInt(key string) int { val, ok := p.data[key].(int) if !ok { return 0 } return val }
The good: It fixed the type assertions and now has a checks, neat!
The bad: The explanation it gave on why it added those doesn't make sense. It didn't actually change those to return copies of values. Let's do one more optimization pass.
Me: Please further optimize this code
ChatGPT: Here are a few ways the provided implementation of the
ParsedJsonBody
type can be further optimized:
- The
Value()
method currently returnsp
if thekey
argument matchesparsedJsonBodyContextKey
, andnil
otherwise. SinceparsedJsonBodyContextKey
is a global variable that is only used in theParseJsonBody()
middleware function, this check is unnecessary. TheValue()
method can simply returnp
without checking thekey
argument.- The
Get()
,GetString()
, andGetInt()
methods currently check if the specified key exists in thedata
field and return a default value if the key does not exist. This check can be removed, because thedata
field is amap[string]interface{}
and trying to access a non-existent key in a map will return the zero value for the type of the map's value (interface{}
in this case, which isnil
).- The
GetString()
andGetInt()
methods currently use theok
return value of the type assertion to determine whether the value for the specified key is of the expected type. This check can be removed, becauseMe: Please show me an implementation with those points
ChatGPT:
type ParsedJsonBody struct { data map[string]interface{} } func (p *ParsedJsonBody) Value(key interface{}) interface{} { // The key argument is not used, so we can simply return p return p } func (p *ParsedJsonBody) Get(key string) interface{} { // We don't need to check if the key exists, because trying // to access a non-existent key in a map will return the zero // value for the type of the map's value (nil in this case) return p.data[key] } func (p *ParsedJsonBody) GetString(key string) string { // We don't need to check the type of the value, because trying // to assert a non-string value to a string type will return // the zero value for the string type ("" in this case) return p.data[key].(string) } func (p *ParsedJsonBody) GetInt(key string) int { // We don't need to check the type of the value, because trying // to assert a non-int value to an int type will return the // zero value for the int type (0 in this case) return p.data[key].(int) }
It got more compact for sure, but the code is wrong. It got rid of the type checks with the confident explanation that type-asserting a non-int value to an int type will return the zero value for the int type, but this is not correct and will panic. Here's the code in Golang playground to try out: https://play.golang.com/p/8jjBJyrJ1uC
While impressive, at this point it's internal state of my program and what I'm trying to do is a bit too messed up, and the confidently wrong explanation+implementation can be dangerous. For a moment I self-doubted my own knowledge of Golang and had to try this out in the playground to make sure.
The network isn't specifically trained on code, so I guess it just doesn't fully understand the inner-workings of Golang, without providing enough context to it first.
Generating unit tests
Let's do something fun - unit tests (groans)! If it can spit out functional unit tests for stuff I'm too lazy to write, that would already be very helpful. After all, it seems to understand code and is able to give explanations for what's going on.
The code here is a snippet from a JWT token parser. It validates the token, extracts sid/uid scopes (UUID) and checks if those are actually valid UUIDs.
Me: Consider the following Golang code:
// other code here func (c *Client) validateToken(t jwt.Token) (*Token, error) { if err := jwt.Validate(t); err != nil { return nil, err } // parse sid sid, ok := t.Get("sid") if !ok { return nil, domain.ErrTokenInvalidSID } sidUUID, err := uuid.Parse(sid.(string)) if err != nil { return nil, domain.ErrTokenInvalidSID } // parse uid uid, ok := t.Get("uid") if !ok { return nil, domain.ErrTokenInvalidUID } uidUUID, err := uuid.Parse(uid.(string)) if err != nil { return nil, domain.ErrTokenInvalidUID } return &Token{ SessionID: sidUUID, UserID: uidUUID, }, nil } // other code here
Please generate some testcases for when sid or uid are invalid
ChatGPT: Here are some test cases for when the
sid
oruid
values in the token are invalid:func TestValidateTokenInvalidSID(t *testing.T) { // Create a new client client := NewClient("some-secret-key") // Create a test token token := jwt.New() token.Set("uid", uuid.New().String()) // Set the sid value to something invalid token.Set("sid", "invalid-sid") token.Set(jwt.ExpirationKey, time.Now().Add(time.Minute * 5).Unix()) // Sign the token signedToken, err := jwt.Sign(token, jwt.WithKey(client.alg, client.hmacSecret)) if err != nil { t.Errorf("unexpected error signing token: %s", err) } // Validate the token _, err := client.validateToken(signedToken) if err != domain.ErrTokenInvalidSID { t.Errorf("expected error to be %s, got %s", domain.ErrTokenInvalidSID, err) } } func TestValidateTokenInvalidUID(t *testing.T) { // Create a new client client := NewClient("some-secret-key") // Create a test token token := jwt.New() // Set the uid value to something invalid token.Set("uid", "invalid-uid") token.Set("sid", uuid.New().String()) token.Set(jwt.ExpirationKey, time.Now().Add(time.Minute * 5).Unix()) // Sign the token signedToken, err := jwt.Sign(token, jwt.WithKey(client.alg, client.hmacSecret)) if err != nil { t.Errorf("unexpected error signing token: %s", err) } // Validate the token _, err := client.validateToken(signedToken) if err != domain.ErrTokenInvalidUID { t.Errorf("expected error to be %s, got %s", domain.ErrTokenInvalidUID, err) } }
Sweet as. Sadly it's missing a bit of context here too, so I can't copy-paste the test as is. jwt.Sign()
isn't returning a jwt.Token
, but the jwt token in signed form, so it can't get passed to validateToken
, but there's no way that it could have known that, without feeding the the jwt package into it first. Besides that it looks good, and the tests run after fixing that!
And to go full circle, now that we have more unit tests, let's ask it to optimize my validateToken
function as well:
One possible optimization for the
validateToken
function would be to move the code that parses thesid
anduid
values into a separate function, and then call that function fromvalidateToken
instead of duplicating the code. This would make the code easier to read and maintain, and would also avoid repeating the same error-checking code multiple times.Here is an example of how the optimized
validateToken
function could look like:func (c *Client) validateToken(t jwt.Token) (*Token, error) { // Validate the token if err := jwt.Validate(t); err != nil { return nil, err } // Parse the sid and uid values from the token sid, err := c.parseSID(t) if err != nil { return nil, err } uid, err := c.parseUID(t) if err != nil { return nil, err } return &Token{ SessionID: sid, UserID: uid, }, nil } func (c *Client) parseSID(t jwt.Token) (uuid.UUID, error) { // parse sid sid, ok := t.Get("sid") if !ok { return uuid.UUID{}, domain.ErrTokenInvalidSID } sidUUID, err := uuid.Parse(sid.(string)) if err != nil { return uuid.UUID{}, domain.ErrTokenInvalidSID } return sidUUID, nil } func (c *Client) parseUID(t jwt.Token) (uuid.UUID, error) { // parse uid uid, ok := t.Get("uid") if !ok { return uuid.UUID{}, domain.ErrTokenInvalidUID } uidUUID, err := uuid.Parse(uid.(string)) if err != nil { return uuid.UUID{}, domain.ErrTokenInvalidUID } return uidUUID, nil }
I was hoping it'd catch the naked .(string)
type assertion, but code splitting is also nice :D
Although this iteration doesn't really avoid repeating the same error-checking code multiple times like it claims.
Collaborative hacking with an AI buddy
I use Copilot quite a lot so I'm already used to having an AI pair-programming buddy, but I strongly believe that as these tools get smarter, they will drastically change the way how we as engineers work.
While not there yet, AI will be able to generate more efficient code faster than us, spot errors we'd miss, or flat-out generate entire applications based on spec input.
When Stable Diffusion came out, some people integrated it into their photo editors to quickly generate content based on a prompt. I am sure we'll see something like this, but for code, in no time. Imagine reading through some code, hitting an "AI" button, then ask it to refactor a routine, or give an explanation what that code does. The plugin will then send all the necessary context to the neural net and ask it to do the task.
The next generation of AI powered plugins could also do things like quickly provide feedback when it strongly believes a mistake has been made, or when it has a high confidence score for a completion / refactoring.
I think that would be really dang cool!
Top comments (1)
its amazing