Modular•4mo ago

Seeking Advice on Porting Python Classes to Mojo Structs

I still struggle with how to port Python classes to Mojo structs in the current state of Mojo and seek some help. Here I post a Python class and suggest four ways to port it to Mojo that I came up with. All of them work in this simple example but seem to have some major drawbacks when the class is more complex. It would be great to get feedback on how you handle this. Here an example Python class to port to Mojo

class Animal:
    def __init__(self, name):
        self.name = name  
    def whoareyou(self):
        return f"I am {self.name}!"
    def speak(self):
        raise NotImplementedError("Subclass must implement abstract method")

class Dog(Animal):
    def speak(self):
        return f"{self.name} says Woof!"

class Cat(Animal):
    def speak(self):
        return f"{self.name} says Meow!"

dog = Dog("Buddy")
cat = Cat("Whiskers")

print(dog.whoareyou()) # Output: I am Buddy!
print(cat.whoareyou()) # Output: I am Whiskers!
print(dog.speak())  # Output: Buddy says Woof!
print(cat.speak())  # Output: Whiskers says Meow!

class Animal:
    def __init__(self, name):
        self.name = name  
    def whoareyou(self):
        return f"I am {self.name}!"
    def speak(self):
        raise NotImplementedError("Subclass must implement abstract method")

class Dog(Animal):
    def speak(self):
        return f"{self.name} says Woof!"

class Cat(Animal):
    def speak(self):
        return f"{self.name} says Meow!"

dog = Dog("Buddy")
cat = Cat("Whiskers")

print(dog.whoareyou()) # Output: I am Buddy!
print(cat.whoareyou()) # Output: I am Whiskers!
print(dog.speak())  # Output: Buddy says Woof!
print(cat.speak())  # Output: Whiskers says Meow!

4 Replies

Martin Dudek•4mo ago

Method 1 My first approach would be to forget about superclasses and implement each subclass independently as struct. The obvious drawback is the redundant code.

struct Dog:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Woof!"

    fn whoareyou(self) -> String:
        return "I am " + self.name + "!" 
    
struct Cat:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Meow!"
    
    fn whoareyou(self) -> String:
        return "I am " + self.name + "!" 
        
fn main():
    var dog = Dog("Buddy")
    var cat = Cat("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

struct Dog:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Woof!"

    fn whoareyou(self) -> String:
        return "I am " + self.name + "!" 
    
struct Cat:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Meow!"
    
    fn whoareyou(self) -> String:
        return "I am " + self.name + "!" 
        
fn main():
    var dog = Dog("Buddy")
    var cat = Cat("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

` Method 2 To avoid the redundant code in Method 1, we could define a trait and an implementation of it for the redundant code.

trait WhoAreYou:
    @staticmethod
    fn whoareyou(name:String) -> String:
        ...

struct IntroduceYourself(WhoAreYou):
    @staticmethod
    fn whoareyou(name:String) -> String:
        return "I am " + name + "!" 

struct Dog[WAY:WhoAreYou]:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Woof!"

    fn whoareyou(self) -> String:
        return WAY.whoareyou(self.name)
    
struct Cat[WAY:WhoAreYou]:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Meow!"
    
    fn whoareyou(self) -> String:
        return WAY.whoareyou(self.name)
        
fn main():
    var dog = Dog[IntroduceYourself]("Buddy")
    var cat = Cat[IntroduceYourself]("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

trait WhoAreYou:
    @staticmethod
    fn whoareyou(name:String) -> String:
        ...

struct IntroduceYourself(WhoAreYou):
    @staticmethod
    fn whoareyou(name:String) -> String:
        return "I am " + name + "!" 

struct Dog[WAY:WhoAreYou]:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Woof!"

    fn whoareyou(self) -> String:
        return WAY.whoareyou(self.name)
    
struct Cat[WAY:WhoAreYou]:
    var name:String
   
    fn __init__(inout self,name:String):
        self.name = name

    fn speak(self) -> String:
        return self.name + " says Meow!"
    
    fn whoareyou(self) -> String:
        return WAY.whoareyou(self.name)
        
fn main():
    var dog = Dog[IntroduceYourself]("Buddy")
    var cat = Cat[IntroduceYourself]("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

` Method 3 This would come closest to subclasses in my opinion, by having a trait for the methods which are implemented individually by the subclasses, and injecting them at compile time. One drawback would be that each method call must transfer all the needed data. (inout ...)

trait Speaker:
    fn __init__(inout self):
        ...
    fn speak(self,name:String) -> String:
        ...

struct Animal[SP:Speaker]:
    var name:String
    var speaker:SP

    fn __init__(inout self,name:String):
        self.name = name
        self.speaker = SP()

    fn speak(self) -> String:
        return self.speaker.speak(self.name)
    
    fn whoareyou(self) -> String:
        return "I am " + self.name + "!" 

struct Dog(Speaker):
    fn __init__(inout self):
        pass
    fn speak(self,name:String) -> String:
        return name + " says Woof!"

struct Cat(Speaker):
    fn __init__(inout self):
        pass
    fn speak(self,name:String) -> String:
        return name + " says  Meow!"
    
fn main():
    var dog = Animal[Dog]("Buddy")
    var cat = Animal[Cat]("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

trait Speaker:
    fn __init__(inout self):
        ...
    fn speak(self,name:String) -> String:
        ...

struct Animal[SP:Speaker]:
    var name:String
    var speaker:SP

    fn __init__(inout self,name:String):
        self.name = name
        self.speaker = SP()

    fn speak(self) -> String:
        return self.speaker.speak(self.name)
    
    fn whoareyou(self) -> String:
        return "I am " + self.name + "!" 

struct Dog(Speaker):
    fn __init__(inout self):
        pass
    fn speak(self,name:String) -> String:
        return name + " says Woof!"

struct Cat(Speaker):
    fn __init__(inout self):
        pass
    fn speak(self,name:String) -> String:
        return name + " says  Meow!"
    
fn main():
    var dog = Animal[Dog]("Buddy")
    var cat = Animal[Cat]("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

` Method 4 Similar to Method 3 but holding the data in the individual structs.

trait Speaker:
    fn __init__(inout self,name:String):
        ...
    fn speak(self) -> String:
        ...
    fn get_name(self)->String:
        ...

struct Animal[SP:Speaker]:
    var speaker:SP

    fn __init__(inout self,name:String):
        self.speaker = SP(name)

    fn speak(self) -> String:
        return self.speaker.speak()
    
    fn whoareyou(self) -> String:
        return "I am " + self.speaker.get_name() + "!" 

struct Dog(Speaker):
    var name:String
    fn __init__(inout self,name:String):
        self.name = name
    fn speak(self) -> String:
        return self.name + " says Woof!"
    fn get_name(self) -> String:
        return self.name

struct Cat(Speaker):
    var name:String
    fn __init__(inout self,name:String):
        self.name = name
    fn speak(self) -> String:
        return self.name + " says  Meow!"
    fn get_name(self) -> String:
        return self.name
    
fn main():
    var dog = Animal[Dog]("Buddy")
    var cat = Animal[Cat]("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

trait Speaker:
    fn __init__(inout self,name:String):
        ...
    fn speak(self) -> String:
        ...
    fn get_name(self)->String:
        ...

struct Animal[SP:Speaker]:
    var speaker:SP

    fn __init__(inout self,name:String):
        self.speaker = SP(name)

    fn speak(self) -> String:
        return self.speaker.speak()
    
    fn whoareyou(self) -> String:
        return "I am " + self.speaker.get_name() + "!" 

struct Dog(Speaker):
    var name:String
    fn __init__(inout self,name:String):
        self.name = name
    fn speak(self) -> String:
        return self.name + " says Woof!"
    fn get_name(self) -> String:
        return self.name

struct Cat(Speaker):
    var name:String
    fn __init__(inout self,name:String):
        self.name = name
    fn speak(self) -> String:
        return self.name + " says  Meow!"
    fn get_name(self) -> String:
        return self.name
    
fn main():
    var dog = Animal[Dog]("Buddy")
    var cat = Animal[Cat]("Whiskers")

    print(dog.whoareyou())
    print(cat.whoareyou())
    print(dog.speak())
    print(cat.speak())

All of these methods work, but I am not really satisfied with any of them. Right now, I'm leaning towards Method 1, as embarrassing as that might be 😜. I understand that Mojo will have the required language features at one point. Big question for me right now, is it too early to port projects from Python which make use of classes ...

Ethan•4mo ago

The proper solution.

trait Namable:
  fn get_name(self) -> String: pass

fn whoareyou[T:Namable](self: T) -> String:
  return str("I am ") + self.get_name() + str("!")

trait Speakable:
  fn speak(self) -> String: pass

trait Introducible:
  fn whoareyou(self) -> String: pass

trait Animal(Namable, Speakable, Introducible): pass

struct Dog(Animal):
  var name: String
  fn __init__(inout self, name: String):
    self.name = name
  
  fn speak(self)-> String:
    return self.name + " says Woof!"
  
  fn whoareyou(self) -> String:
      return whoareyou(self)
  
  fn get_name(self) -> String:
    return self.name
  
struct Cat(Animal):
  var name: String
  fn __init__(inout self, name: String):
    self.name = name
  
  fn speak(self)-> String:
    return self.name + " says Meow!"
  
  fn whoareyou(self) -> String:
      return whoareyou(self)
  
  fn get_name(self) -> String:
    return self.name

fn main():
  var dog = Dog("Buddy")
  var cat = Cat("Whiskers")

  print(dog.whoareyou())
  print(dog.speak())
  print(cat.whoareyou())
  print(cat.speak())

trait Namable:
  fn get_name(self) -> String: pass

fn whoareyou[T:Namable](self: T) -> String:
  return str("I am ") + self.get_name() + str("!")

trait Speakable:
  fn speak(self) -> String: pass

trait Introducible:
  fn whoareyou(self) -> String: pass

trait Animal(Namable, Speakable, Introducible): pass

struct Dog(Animal):
  var name: String
  fn __init__(inout self, name: String):
    self.name = name
  
  fn speak(self)-> String:
    return self.name + " says Woof!"
  
  fn whoareyou(self) -> String:
      return whoareyou(self)
  
  fn get_name(self) -> String:
    return self.name
  
struct Cat(Animal):
  var name: String
  fn __init__(inout self, name: String):
    self.name = name
  
  fn speak(self)-> String:
    return self.name + " says Meow!"
  
  fn whoareyou(self) -> String:
      return whoareyou(self)
  
  fn get_name(self) -> String:
    return self.name

fn main():
  var dog = Dog("Buddy")
  var cat = Cat("Whiskers")

  print(dog.whoareyou())
  print(dog.speak())
  print(cat.whoareyou())
  print(cat.speak())

In the future when trait supports default methods and fields the following would just work:

trait Namable:
  var name: String

fn whoareyou[T:Namable](self: T) -> String:
  return str("I am ") + self.name + str("!")

trait Speakable:
  fn speak(self) -> String: pass

trait Introducible:
  fn whoareyou(self) -> String: 
    return whoareyou(self)

trait Animal(Namable, Speakable, Introducible):
  fn __init__(inout self, name: String):
    self.name = name

struct Dog(Animal):
  fn speak(self)-> String:
    return self.name + " says Woof!"
  
struct Cat(Animal):
  fn speak(self)-> String:
    return self.name + " says Meow!"

fn main():
  var dog = Dog("Buddy")
  var cat = Cat("Whiskers")

  print(dog.whoareyou())
  print(dog.speak())
  print(cat.whoareyou())
  print(cat.speak())

trait Namable:
  var name: String

fn whoareyou[T:Namable](self: T) -> String:
  return str("I am ") + self.name + str("!")

trait Speakable:
  fn speak(self) -> String: pass

trait Introducible:
  fn whoareyou(self) -> String: 
    return whoareyou(self)

trait Animal(Namable, Speakable, Introducible):
  fn __init__(inout self, name: String):
    self.name = name

struct Dog(Animal):
  fn speak(self)-> String:
    return self.name + " says Woof!"
  
struct Cat(Animal):
  fn speak(self)-> String:
    return self.name + " says Meow!"

fn main():
  var dog = Dog("Buddy")
  var cat = Cat("Whiskers")

  print(dog.whoareyou())
  print(dog.speak())
  print(cat.whoareyou())
  print(cat.speak())

Martin Dudek•4mo ago

Thanks a lot @Ethan , this is as very elagant way of implementing it, very helpful. Yes fields and default methods for traits would enable more straigt forward implementations, but your solution is for what we can do now the best i saw so far. Thanks again

Ethan•4mo ago

My pleasure. This is the de facto way to achieve polymorphism with traits. I am glad you got the hang of it.

Gaming

Programming

Seeking Advice on Porting Python Classes to Mojo Structs